open Printf
open Str
open Lexing
open Types

(** Empty location, use when node location is unkown or irrelevant. *)
let noloc = ("", 0, 0, 0, 0)

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_node node = prerr_endline (Stringify.node2str node)

let log_plain_line verbosity line =
  if Globals.args.verbose >= verbosity then prerr_endline line

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

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

(* Variable generation *)

let generate_id id num = "_" ^ id ^ "_" ^ string_of_int num

let var_counter = ref 0

let fresh_id base =
  (* For generated variables, just replace the counter *)
  let base =
    if string_match (regexp "^_\\(.+\\)_[0-9]+_?$") base 0 then
      matched_group 1 base
    else
      base
  in
  var_counter := !var_counter + 1;
  generate_id base !var_counter

(* Constants are marked by a leading _ for recognition during constant
 * propagation *)
let generate_const id num = generate_id id num ^ "_"

let fresh_const id = fresh_id id ^ "_"

let is_generated_id id = String.length id >= 1 & id.[0] = '_'

let is_const_id id =
  String.length id >= 2
  & id.[0] = '_'
  & id.[String.length id - 1] = '_'

let loc_from_lexpos pstart pend =
  let (fname, ystart, yend, xstart, xend) = begin
    pstart.pos_fname,
    pstart.pos_lnum,
    pend.pos_lnum,
    (pstart.pos_cnum - pstart.pos_bol + 1),
    (pend.pos_cnum - pend.pos_bol)
  end 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 rec traverse u ( *: ) trav node =
  let trav_all nodes =
    let (nodes, res) = List.split (List.map trav nodes) in
    (nodes, List.fold_left ( *: ) u res)
  in
  match node with
  | Program (decls, ann) ->
    let (decls, res_decls) = trav_all decls in
    (Program (flatten_blocks decls, ann), res_decls)
  | FunDec (ret_type, name, params, ann) ->
    let (params, res_params) = trav_all params in
    (FunDec (ret_type, name, params, ann), res_params)
  | FunDef (export, ret_type, name, params, body, ann) ->
    let (params, resp) = trav_all params in
    let (body, resb) = trav body in
    (FunDef (export, ret_type, name, params, body, ann), resp *: resb)
  | GlobalDec (ctype, name, ann) ->
    (GlobalDec (ctype, name, ann), u)
  | GlobalDef (export, ctype, name, Some init, ann) ->
    let (init, res_init) = trav init in
    (GlobalDef (export, ctype, name, Some init, ann), res_init)

  | VarDecs decs ->
    let (decs, res_decs) = trav_all decs in
    (VarDecs decs, res_decs)
  | LocalFuns funs ->
    let (funs, res_funs) = trav_all funs in
    (LocalFuns funs, res_funs)

  | VarDec (ctype, name, Some init, ann) ->
    let (init, res_init) = trav init in
    (VarDec (ctype, name, Some init, ann), res_init)
  | Assign (name, None, value, ann) ->
    let (value, res_value) = trav value in
    (Assign (name, None, value, ann), res_value)
  | Assign (name, Some dims, value, ann) ->
    let (dims, res_dims) = trav_all dims in
    let (value, res_value) = trav value in
    (Assign (name, Some dims, value, ann), res_dims *: res_value)
  | VarLet (dec, None, value, ann) ->
    let (value, res_value) = trav value in
    (VarLet (dec, None, value, ann), res_value)
  | VarLet (dec, Some dims, value, ann) ->
    let (dims, res_dims) = trav_all dims in
    let (value, res_value) = trav value in
    (VarLet (dec, Some dims, value, ann), res_dims *: res_value)
  | Return (value, ann) ->
    let (value, res_value) = trav value in
    (Return (value, ann), res_value)
  | If (cond, body, ann) ->
    let (cond, res_cond) = trav cond in
    let (body, res_body) = trav body in
    (If (cond, body, ann), res_cond *: res_body)
  | IfElse (cond, tbody, fbody, ann) ->
    let (cond, resa) = trav cond in
    let (tbody, resb) = trav tbody in
    let (fbody, resc) = trav fbody in
    (IfElse (cond, tbody, fbody, ann), resa *: resb *: resc)
  | While (cond, body, ann) ->
    let (cond, resc) = trav cond in
    let (body, resb) = trav body in
    (While (cond, body, ann), resc *: resb)
  | DoWhile (cond, body, ann) ->
    let (cond, resc) = trav cond in
    let (body, resb) = trav body in
    (DoWhile (cond, body, ann), resc *: resb)
  | For (counter, start, stop, step, body, ann) ->
    let (start, resa) = trav start in
    let (stop, resb) = trav stop in
    let (step, resc) = trav step in
    let (body, resd) = trav body in
    let res = resa *: resb *: resc *: resd in
    (For (counter, start, stop, step, body, ann), res)
  | Allocate (dec, dims, ann) ->
    let (dims, res_dims) = trav_all dims in
    (Allocate (dec, dims, ann), res_dims)
  | Expr value ->
    let (value, res_value) = trav value in
    (Expr value, res_value)
  | Block (body) ->
    let (body, res_body) = trav_all body in
    (Block body, res_body)

  | Monop (op, value, ann) ->
    let (value, res_value) = trav value in
    (Monop (op, value, ann), res_value)
  | Binop (op, left, right, ann) ->
    let (left, res_left) = trav left in
    let (right, res_right) = trav right in
    (Binop (op, left, right, ann), res_left *: res_right)
  | Cond (cond, texpr, fexpr, ann) ->
    let (cond, resa) = trav cond in
    let (texpr, resb) = trav texpr in
    let (fexpr, resc) = trav fexpr in
    (Cond (cond, texpr, fexpr, ann), resa *: resb *: resc)
  | TypeCast (ctype, value, ann) ->
    let (value, res_value) = trav value in
    (TypeCast (ctype, value, ann), res_value)
  | FunCall (name, args, ann) ->
    let (args, res_args) = trav_all args in
    (FunCall (name, args, ann), res_args)
  | Arg value ->
    let (value, res_value) = trav value in
    (Arg value, res_value)

  | ArrayInit (value, dims) ->
    let (value, res_value) = trav value in
    (ArrayInit (value, dims), res_value)
  | Var (dec, Some dims, ann) ->
    let (dims, res_dims) = trav_all dims in
    (Var (dec, Some dims, ann), res_dims)
  | VarUse (dec, Some dims, ann) ->
    let (dims, res_dims) = trav_all dims in
    (VarUse (dec, Some dims, ann), res_dims)
  | FunUse (dec, params, ann) ->
    let (params, res_params) = trav_all params in
    (FunUse (dec, params, ann), res_params)

  | _ -> (node, u)

let traverse_unit visit node =
  fst (traverse () (fun () () -> ()) (fun n -> (visit n, ())) node)

let traverse_list visit = traverse [] (@) visit

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, _)
  | 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 count_tabs str upto =
  let rec count n = function
    | 0 -> n
    | i -> count (if String.get str (i - 1) = '\t' then n + 1 else n) (i - 1)
  in count 0 upto

let tabwidth = 4

let retab str = global_replace (regexp "\t") (repeat " " tabwidth) str

let indent n = repeat (repeat " " (tabwidth - 1)) n

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 begin
      prerr_endline (retab line);
      prerr_string (indent (count_tabs line right));
      for i = 1 to left - 1 do prerr_char ' ' done;
      for i = left to right do prerr_char '^' done;
      prerr_endline "";
    end
  done;
  ()

let prerr_loc_msg loc msg =
  if Globals.args.verbose >= 1 then begin
    let (fname, ystart, yend, xstart, xend) = loc in
    if loc != noloc then begin
      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;
    end;
    eprintf "%s\n" msg;

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

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 ->
      let hd = f i hd in
      hd :: (trav (i + 1) tl)
  in trav 0 lst

(** Constants that are *)
let immediate_consts = [
  BoolVal true;
  BoolVal false;
  IntVal (-1l);
  IntVal 0l;
  IntVal 1l;
  FloatVal 0.0;
  FloatVal 1.0;
]

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

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

let node_error node msg =
  prerr_loc_msg (locof node) ("Error: " ^ msg)

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

let print_error = function
  | Msg msg ->
    eprintf "Error: %s\n" msg;

  | LocMsg (loc, msg) ->
    prerr_loc_msg loc ("Error: " ^ msg)

  | NodeMsg (node, msg) ->
    (* If no location is given, just stringify the node to at least give
      * some information *)
    let msg = if locof node = noloc then
      msg ^ "\nnode: " ^ Stringify.node2str node
    else msg in
    node_error node msg

  | NoMsg -> ()

let quit_on_error node = function
  | [] -> node
  | errors ->
    List.iter print_error errors;
    raise (FatalError NoMsg)