open Printf
open Lexing
open Types

let rec repeat s n = if n < 1 then "" else s ^ (repeat s (n - 1))
let expand n text = text ^ repeat " " (n - String.length text)

(* Logging functions *)

let hline = "-----------------------------------------------------------------"

let prt_line = prerr_endline

let prt_node node = prt_line (Stringify.node2str node)

let log_plain_line verbosity line =
    if args.verbose >= verbosity then prt_line line

let log_line verbosity line =
    log_plain_line verbosity (repeat " " 13 ^ line)

let log_node verbosity node =
    if args.verbose >= verbosity then prt_node node

(* Variable generation *)
let var_counter = ref 0
let fresh_var prefix =
    var_counter := !var_counter + 1;
    prefix ^ "$" ^ string_of_int !var_counter

(* Constants are marked by a double $$ for recognition during constant
 * propagation *)
let fresh_const prefix = fresh_var (prefix ^ "$")

let loc_from_lexpos pstart pend =
    let (fname, ystart, yend, xstart, xend) = (
        pstart.pos_fname,
        pstart.pos_lnum,
        pend.pos_lnum,
        (pstart.pos_cnum - pstart.pos_bol + 1),
        (pend.pos_cnum - pend.pos_bol)
    ) in
    if ystart = yend && xend < xstart then
        (fname, ystart, yend, xstart, xstart)
    else
        (fname, ystart, yend, xstart, xend)

let rec flatten_blocks lst =
    let flatten = flatten_blocks in
    let rec trav = function
        | Block body ->
            Block (flatten body)
        | FunDef (export, ret_type, name, params, body, ann) ->
            FunDef (export, ret_type, name, flatten params, trav body, ann)
        | If (cond, body, ann) ->
            If (cond, trav body, ann)
        | IfElse (cond, tbody, fbody, ann) ->
            IfElse (cond, trav tbody, trav fbody, ann)
        | While (cond, body, ann) ->
            While (cond, trav body, ann)
        | DoWhile (cond, body, ann) ->
            DoWhile (cond, trav body, ann)
        | For (counter, start, stop, step, body, ann) ->
            For (counter, start, stop, step, trav body, ann)
        | VarDecs decs ->
            VarDecs (flatten decs)
        | LocalFuns decs ->
            LocalFuns (flatten decs)
        | node -> node
    in
    match lst with
    | []                -> []
    | Block nodes :: tl -> flatten nodes @ (flatten tl)
    | DummyNode :: tl   -> flatten tl
    | hd :: tl          -> trav hd :: (flatten tl)

(* Default tree transformation
 * (node -> node) -> node -> node *)
let transform_children trav node =
    let trav_all nodes = List.map trav nodes in
    match node with
    | Program (decls, ann) ->
        Program (flatten_blocks (trav_all decls), ann)
    | FunDec (ret_type, name, params, ann) ->
        FunDec (ret_type, name, trav_all params, ann)
    | FunDef (export, ret_type, name, params, body, ann) ->
        FunDef (export, ret_type, name, trav_all params, trav body, ann)
    | GlobalDec (ctype, name, ann) ->
        GlobalDec (ctype, name, ann)
    | GlobalDef (export, ctype, name, Some init, ann) ->
        GlobalDef (export, ctype, name, Some (trav init), ann)

    | VarDecs decs ->
        VarDecs (trav_all decs)
    | LocalFuns funs ->
        LocalFuns (trav_all funs)

    | VarDec (ctype, name, Some init, ann) ->
        VarDec (ctype, name, Some (trav init), ann)
    | Assign (name, None, value, ann) ->
        Assign (name, None, trav value, ann)
    | Assign (name, Some dims, value, ann) ->
        Assign (name, Some (trav_all dims), trav value, ann)
    | VarLet (dec, None, value, ann) ->
        VarLet (dec, None, trav value, ann)
    | VarLet (dec, Some dims, value, ann) ->
        VarLet (dec, Some (trav_all dims), trav value, ann)
    | Return (value, ann) ->
        Return (trav value, ann)
    | If (cond, body, ann) ->
        If (trav cond, trav body, ann)
    | IfElse (cond, true_body, false_body, ann) ->
        IfElse (trav cond, trav true_body, trav false_body, ann)
    | While (cond, body, ann) ->
        While (trav cond, trav body, ann)
    | DoWhile (cond, body, ann) ->
        DoWhile (trav cond, trav body, ann)
    | For (counter, start, stop, step, body, ann) ->
        For (counter, trav start, trav stop, trav step, trav body, ann)
    | Allocate (dec, dims, ann) ->
        Allocate (dec, trav_all dims, ann)
    | Expr value ->
        Expr (trav value)
    | Block (body) ->
        Block (trav_all body)

    | Monop (op, value, ann) ->
        Monop (op, trav value, ann)
    | Binop (op, left, right, ann) ->
        Binop (op, trav left, trav right, ann)
    | Cond (cond, true_expr, false_expr, ann) ->
        Cond (trav cond, trav true_expr, trav false_expr, ann)
    | TypeCast (ctype, value, ann) ->
        TypeCast (ctype, trav value, ann)
    | FunCall (name, args, ann) ->
        FunCall (name, trav_all args, ann)
    | Arg value ->
        Arg (trav value)

    | ArrayInit (value, dims) ->
        ArrayInit (trav value, dims)
    | ArrayScalar value ->
        ArrayScalar (trav value)
    | Var (dec, Some dims, ann) ->
        Var (dec, Some (trav_all dims), ann)
    | VarUse (dec, Some dims, ann) ->
        VarUse (dec, Some (trav_all dims), ann)
    | FunUse (dec, params, ann) ->
        FunUse (dec, trav_all params, ann)

    | _ -> node

let annotate a = function
    | Program (decls, ann) ->
        Program (decls, a :: ann)
    | FunDec (ret_type, name, params, ann) ->
        FunDec (ret_type, name, params, a :: ann)
    | FunDef (export, ret_type, name, params, body, ann) ->
        FunDef (export, ret_type, name, params, body, a :: ann)
    | GlobalDec (ctype, name, ann) ->
        GlobalDec (ctype, name, a :: ann)
    | GlobalDef (export, ctype, name, init, ann) ->
        GlobalDef (export, ctype, name, init, a :: ann)
    | VarDec (ctype, name, init, ann) ->
        VarDec (ctype, name, init, a :: ann)
    | Assign (name, dims, value, ann) ->
        Assign (name, dims, value, a :: ann)
    | VarLet (dec, dims, value, ann) ->
        VarLet (dec, dims, value, a :: ann)
    | Return (value, ann) ->
        Return (value, a :: ann)
    | If (cond, body, ann) ->
        If (cond, body, a :: ann)
    | IfElse (cond, true_body, false_body, ann) ->
        IfElse (cond, true_body, false_body, a :: ann)
    | While (cond, body, ann) ->
        While (cond, body, a :: ann)
    | DoWhile (cond, body, ann) ->
        DoWhile (cond, body, a :: ann)
    | For (counter, start, stop, step, body, ann) ->
        For (counter, start, stop, step, body, a :: ann)
    | Allocate (dec, dims, ann) ->
        Allocate (dec, dims, a :: ann)
    | Monop (op, value, ann) ->
        Monop (op, value, a :: ann)
    | Binop (op, left, right, ann) ->
        Binop (op, left, right, a :: ann)
    | Cond (cond, true_expr, false_expr, ann) ->
        Cond (cond, true_expr, false_expr, a :: ann)
    | TypeCast (ctype, value, ann) ->
        TypeCast (ctype, value, a :: ann)
    | FunCall (name, args, ann) ->
        FunCall (name, args, a :: ann)
    | Arg value ->
        Arg (value)
    | Var (dec, dims, ann) ->
        Var (dec, dims, a :: ann)
    | VarUse (dec, dims, ann) ->
        VarUse (dec, dims, a :: ann)
    | FunUse (dec, params, ann) ->
        FunUse (dec, params, a :: ann)
    | Const (BoolVal value, ann) ->
        Const (BoolVal value, a :: ann)
    | Const (IntVal value, ann) ->
        Const (IntVal value, a :: ann)
    | Const (FloatVal value, ann) ->
        Const (FloatVal value, a :: ann)
    | ArrayConst (value, ann) ->
        ArrayConst (value, a :: ann)
    | Param (ctype, name, ann) ->
        Param (ctype, name, a :: ann)
    | Dim (name, ann) ->
        Dim (name, a :: ann)

    | _ -> raise InvalidNode

let rec annof = function
    | Program (_, ann)
    | Param (_, _, ann)
    | Dim (_, ann)
    | FunDec (_, _, _, ann)
    | FunDef (_, _, _, _, _, ann)
    | GlobalDec (_, _, ann)
    | GlobalDef (_, _, _, _, ann)
    | VarDec (_, _, _, ann)
    | Assign (_, _, _, ann)
    | VarLet (_, _, _, ann)
    | Return (_, ann)
    | If (_, _, ann)
    | IfElse (_, _, _, ann)
    | While (_, _, ann)
    | DoWhile (_, _, ann)
    | For (_, _, _, _, _, ann)
    | Allocate (_, _, ann)
    | Const (BoolVal _, ann)
    | Const (IntVal _, ann)
    | Const (FloatVal _, ann)
    | ArrayConst (_, ann)
    | Var (_, _, ann)
    | Monop (_, _, ann)
    | Binop (_, _, _, ann)
    | Cond (_, _, _, ann)
    | TypeCast (_, _, ann)
    | VarUse (_, _, ann)
    | FunUse (_, _, ann)
    | FunCall (_, _, ann) -> ann

    | ArrayInit (value, _)
    | ArrayScalar value
    | Expr value
    | Arg value -> annof value

    | _ -> raise InvalidNode

let locof node =
     let rec trav = function
         | []           -> noloc
         | Loc loc :: _ -> loc
         | _ :: tl      -> trav tl
     in trav (annof node)

let depthof node =
     let rec trav = function
         | [] ->
            prerr_string "cannot get depth for: ";
            prt_node node;
            raise InvalidNode
         | Depth depth :: _ -> depth
         | _ :: tl          -> trav tl
     in trav (annof node)

let indexof node =
     let rec trav = function
         | [] ->
            prerr_string "cannot get index for: ";
            prt_node node;
            raise InvalidNode
         | Index index :: _ -> index
         | _ :: tl          -> trav tl
     in trav (annof node)

let typeof = function
    (* Some nodes have their type as property *)
    | VarDec (ctype, _, _, _)
    | Param (ctype, _, _)
    | FunDec (ctype, _, _, _)
    | FunDef (_, ctype, _, _, _, _)
    | GlobalDec (ctype, _, _)
    | GlobalDef (_, ctype, _, _, _)
    | TypeCast (ctype, _, _)
        -> ctype

    (* Dim nodes are always type Int, and are copied by context analysis before
     * they are annotated with Type Int, so this match is necessary *)
    | Dim _ -> Int

    (* Other nodes must be annotated during typechecking *)
    | node ->
        let rec trav = function
         | [] ->
            prerr_string "cannot get type for: ";
            prt_node node;
            raise InvalidNode
            | Type t :: _ -> t
            | _ :: tl     -> trav tl
        in trav (annof node)

let labelof node =
     let rec trav = function
         | [] ->
            prerr_string "cannot get label for: ";
            prt_node node;
            raise InvalidNode
         | LabelName label :: _ -> label
         | _ :: tl              -> trav tl
     in trav (annof node)

let const_type = function
    | BoolVal _  -> Bool
    | IntVal _   -> Int
    | FloatVal _ -> Float

(*
let get_line str n =
    let rec find_start from = function
        | n when n < 1 -> raise (Invalid_argument "n")
        | 1 -> from
        | n -> find_start ((String.index_from str from '\n') + 1) (n - 1)
    in
    let linestart = find_start 0 n in
    let len = String.length str in
    let lineend =
        try String.index_from str linestart '\n'
        with Not_found -> len
    in
    String.sub str linestart (lineend - linestart)
*)

let prerr_loc (fname, ystart, yend, xstart, xend) =
    let file = open_in fname in

    (* skip lines until the first matched line *)
    for i = 1 to ystart - 1 do let _ = input_line file in () done;

    (* for each line in `loc`, print the source line with an underline *)
    for l = ystart to yend do
        let line = input_line file in
        let linewidth = String.length line in
        let left = if l = ystart then xstart else 1 in
        let right = if l = yend then xend else linewidth in
        if linewidth > 0 then (
            prerr_endline line;
            for i = 1 to left - 1 do prerr_char ' ' done;
            for i = left to right do prerr_char '^' done;
            prerr_endline "";
        )
    done;
    ()

let prerr_loc_msg loc msg =
    if args.verbose >= 1 then (
        let (fname, ystart, yend, xstart, xend) = loc in
        if loc != noloc then (
            let line_s = if yend != ystart
                then sprintf "lines %d-%d" ystart yend
                else sprintf "line %d" ystart
            in
            let char_s = if xend != xstart || yend != ystart
                then sprintf "characters %d-%d" xstart xend
                else sprintf "character %d" xstart
            in
            eprintf "File \"%s\", %s, %s:\n" fname line_s char_s;
        );
        eprintf "%s\n" msg;

        if args.verbose >= 1 && loc != noloc then
            try prerr_loc loc
            with Sys_error _ -> ()
    );
    ()

let block_body = function
    | Block nodes -> nodes
    | _ -> raise InvalidNode

let basetypeof node = match typeof node with
    | ArrayDims (ctype, _)
    | Array ctype
    | ctype -> ctype

let nameof = function
    | GlobalDec (_, name, _)
    | GlobalDef (_, _, name, _, _)
    | FunDec (_, name, _, _)
    | FunDef (_, _, name, _, _, _)
    | VarDec (_, name, _, _)
    | Param (_, name, _)
    | Dim (name, _) -> name
    | _ -> raise InvalidNode

let optmap f = function
    | None -> None
    | Some lst -> Some (List.map f lst)

let optmapl f = function
    | None -> []
    | Some lst -> List.map f lst

let mapi f lst =
    let rec trav i = function
        | [] -> []
        | hd :: tl -> f i hd :: (trav (i + 1) tl)
    in trav 0 lst

let is_immediate_const const =
    if args.optimize then List.mem const immediate_consts else false

let is_array node = match typeof node with
    | ArrayDims _ | Array _ -> true
    | _ -> false

let node_warning node msg =
    prerr_loc_msg (locof node) ("Warning: " ^ msg)