open Printf
open Lexing
open Ast

let var_counter = ref 0

let fresh_var prefix =
    var_counter := !var_counter + 1;
    prefix ^ "$" ^ string_of_int !var_counter

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)

(* 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, loc) ->
        Program (trav_all decls, loc)
    | FunDec (ret_type, name, params, loc) ->
        FunDec (ret_type, name, trav_all params, loc)
    | FunDef (export, ret_type, name, params, body, loc) ->
        FunDef (export, ret_type, name, trav_all params, trav body, loc)
    | GlobalDec (ctype, name, loc) ->
        GlobalDec (ctype, name, loc)
    | GlobalDef (export, ctype, name, Some init, loc) ->
        GlobalDef (export, ctype, name, Some (trav init), loc)

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

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

    | ArrayScalar (value, dims) ->
        ArrayScalar (trav value, dims)
    | Type (value, ctype) ->
        Type (trav value, ctype)
    | VarLet (assign, def, depth) ->
        VarLet (trav assign, def, depth)
    | VarUse (var, def, depth) ->
        VarUse (trav var, def, depth)
    | FunUse (funcall, def, depth) ->
        FunUse (trav funcall, def, depth)
    | DimDec node ->
        DimDec (trav node)

    | _ -> node

(* Default tree transformation
 * (node -> node) -> node -> node *)
let rec transform_all trav = function
    | [] -> []
    | node :: tail -> trav node :: (transform_all trav tail)

 let rec locof = function
    | Program (_, loc)
    | Param (_, _, loc)
    | Dim (_, loc)
    | FunDec (_, _, _, loc)
    | FunDef (_, _, _, _, _, loc)
    | GlobalDec (_, _, loc)
    | GlobalDef (_, _, _, _, loc)
    | VarDec (_, _, _, loc)
    | Assign (_, _, _, loc)
    | Return (_, loc)
    | If (_, _, loc)
    | IfElse (_, _, _, loc)
    | While (_, _, loc)
    | DoWhile (_, _, loc)
    | For (_, _, _, _, _, loc)
    | Allocate (_, _, _, loc)
    | BoolConst (_, loc)
    | IntConst (_, loc)
    | FloatConst (_, loc)
    | ArrayConst (_, loc)
    | Var (_, loc)
    | Deref (_, _, loc)
    | Monop (_, _, loc)
    | Binop (_, _, _, loc)
    | Cond (_, _, _, loc)
    | TypeCast (_, _, loc)
    | FunCall (_, _, loc) -> loc

    | ArrayScalar (value, _)
    | Expr value
    | VarLet (value, _, _)
    | VarUse (value, _, _)
    | FunUse (value, _, _)
    | Arg value
    | Type (value, _)
    | DimDec value -> locof value

    | _ -> noloc

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 input_line file 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 verbose =
    let (fname, ystart, yend, xstart, xend) = loc in
    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 verbose >= 2 then prerr_loc loc;
    ()

let rec flatten_blocks = function
    | [] -> []
    | Block nodes :: t -> (flatten_blocks nodes) @ (flatten_blocks t)
    | h :: t -> h :: (flatten_blocks t)

let ctypeof = function
    | VarDec (ctype, _, _, _)
    | Param (ctype, _, _)
    | FunDec (ctype, _, _, _)
    | FunDef (_, ctype, _, _, _, _)
    | GlobalDec (ctype, _, _)
    | GlobalDef (_, ctype, _, _, _)
    | TypeCast (ctype, _, _)
    | Type (_, ctype)
        -> ctype

    | DimDec _ -> Int

    | _ -> raise InvalidNode

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

let rec list_size = function
    | [] -> 0
    | hd :: tl -> 1 + (list_size tl)

let base_type = function
    | ArrayDec (ctype, _)
    | ArrayDef (ctype, _)
    | ctype -> ctype