open Printf
open Types
open Util
open Stringify

let comline comment = InlineComment (EmptyLine, comment)

let assemble program =
  let labcounter = ref 0 in
  let genlabel suffix =
  labcounter := !labcounter + 1;
  string_of_int !labcounter ^ "_" ^ suffix
  in

  let consts = Hashtbl.create 20 in

  let rec trav node =
    let rec trav_all = function
      | [] -> []
      | hd :: tl -> trav hd @ (trav_all tl)
    in
    let rec traverse_localfuns = function
      | LocalFuns body -> List.concat (List.map trav body)
      | Block body -> List.concat (List.map traverse_localfuns body)
      | _ -> []
    in
    match node with
    (* Global *)

    | Program (decls, _) ->
      trav_all decls

    | GlobalDef (_, ctype, name, _, _) ->
      [Comment (sprintf "global var \"%s\" at index %d" name (indexof node));
       Global ctype]

    | FunDec (ret_type, name, params, _) ->
      [Comment (sprintf "extern fun \"%s\" at index %d" name (indexof node));
       Import (name, ret_type, List.map typeof params)]

    | FunDef (export, ret_type, name, params, body, _) ->
      let label = labelof node in
      begin
        if export then
          let param_types = List.map typeof params in
          [Export (name, ret_type, param_types, label)]
        else []
      end @
      [Comment (sprintf "function \"%s\" with %d parameters and %d local vars"
                        label (List.length params) (indexof node));
       Label label;
       RtnEnter (indexof node)] @
      (trav_all params) @
      (trav body) @
      (match ret_type with Void -> [Ret Void] | _ -> []) @
      [EmptyLine] @
      (traverse_localfuns body)

    | VarDec (_, name, _, _) ->
      [comline (sprintf "local var \"%s\" at index %d" name (indexof node))]

    | Param (_, name, _) ->
      [comline (sprintf "parameter \"%s\" at index %d" name (indexof node))]

    | LocalFuns _ -> []

    | Block body | VarDecs body -> trav_all body

    (* Statements *)

    | VarLet (dec, None, value, _) ->
      let store = match (depthof dec, depthof node) with
      | (0, _)            -> Store (typeof dec, Glob,        indexof dec)
      | (a, b) when a = b -> Store (typeof dec, Current,     indexof dec)
      | (a, b)            -> Store (typeof dec, Rel (b - a), indexof dec)
      in
      trav value @ [InlineComment (store, node2str node)]

    | Return (value, _) ->
      trav value @ [InlineComment (Ret (typeof value), node2str node)]

    | If (cond, body, _) ->
      let endlabel = genlabel "end" in
      (trav cond) @
      [Branch (false, endlabel);
       comline ("if (" ^ (node2str cond) ^ ") {")] @
      (trav body) @
      [comline "}";
       Label endlabel]

    | IfElse (cond, true_body, false_body, _) ->
      let elselabel = genlabel "else" in
      let endlabel = genlabel "end" in
      (trav cond) @
      [Branch (false, elselabel);
       comline ("if (" ^ (node2str cond) ^ ") {")] @
      (trav true_body) @
      [Jump endlabel;
       comline "} else {";
       Label elselabel] @
      (trav false_body) @
      [comline "}";
       Label endlabel]

    | While (cond, body, _) ->
      let startlabel = genlabel "while" in
      let endlabel = genlabel "end" in
      let com = ("while (" ^ (node2str cond) ^ ") {") in
      [Label startlabel] @
      (trav cond) @
      [InlineComment (Branch (false, endlabel), com)] @
      (trav body) @
      [Jump startlabel;
       Label endlabel;
       comline "}"]

    | DoWhile (cond, body, _) ->
      let startlabel = genlabel "dowhile" in
      let com = ("} while (" ^ (node2str cond) ^ ");") in
      [comline "do {";
       Label startlabel] @
      (trav body) @
      (trav cond) @
      [InlineComment (Branch (true, startlabel), com)]

    (* Expression statement pops the disregarded expression value from the
     * stack, if any *)
    | Expr value ->
      let pop = match typeof value with
      | Void -> []
      | ctype -> [InlineComment (Pop ctype, "disregard return value")]
      in
      trav value @ pop

    (* Expressions *)

    (* Immediate values are handled here, and not in the peephole optimizer, for
     * convenience: the indices in the constant table would be altered, so
     * entries cannot be removed. By this early detection (also during index
     * analysis), they are not added in the first place *)
    | Const (value, _) when is_immediate_const value ->
      [InlineComment (LoadImm value, node2str node)]

    | Const (value, _) ->
      Hashtbl.replace consts value (indexof node);
      let load = LoadConst (typeof node, indexof node) in
      [InlineComment (load, node2str node)]

    | VarUse (dec, None, _) ->
      let load = match (depthof dec, depthof node) with
      | (0, _)            -> Load (typeof dec, Glob,        indexof dec)
      | (a, b) when a = b -> Load (typeof dec, Current,     indexof dec)
      | (a, b)            -> Load (typeof dec, Rel (b - a), indexof dec)
      in
      [InlineComment (load, node2str node)]

    | Monop (op, value, _) ->
      (trav value) @
      [InlineComment (Op (op, typeof value), op2str op)]

    | Binop (op, left, right, _) ->
      (trav left) @
      (trav right) @
      [InlineComment (Op (op, typeof left), op2str op)]

    | TypeCast (ctype, value, _) ->
      let vtype = typeof value in
      begin
        match (ctype, vtype) with
        | (Float, Int) | (Int, Float) -> ()
        | _ ->
          let msg = sprintf
            "invalid typecast: %s -> %s"
            (type2str vtype) (type2str ctype)
          in
          raise (FatalError (NodeMsg (node, msg)))
      end;
      trav value @ [Convert (vtype, ctype)]

    (* Function calls *)
    | FunUse (dec, args, _) ->
      let init = match (depthof dec, depthof node) with
      | (0, _)                -> RtnInit Glob
      | (a, b) when a = b - 1 -> RtnInit Current
      | (a, b) when a = b     -> RtnInit Local
      | (a, b)                -> RtnInit (Rel (b - a - 1))
      in
      let jmp = match dec with
      | FunDec _ -> RtnJmp (ExternFun (indexof dec))
      | FunDef _ -> RtnJmp (LocalFun (List.length args, labelof dec))
      | _ -> raise InvalidNode
      in
      InlineComment (init, nameof dec) ::
      (trav_all args) @
      [InlineComment (jmp, node2str node)]

    | Arg value -> trav value

    (* Conditional expression (short-circuit evaluation) *)
    (*     <cond>
    *     branch_f else
    *     <true_expr>
    *     jump end
    * else:
    *     <false_expr>
    * end:
    *)
    | Cond (cond, texp, fexp, _) ->
      let elselabel = genlabel "false_expr" in
      let endlabel = genlabel "end" in
      (trav cond) @
      [Branch (false, elselabel)] @
      (trav texp) @
      [Jump (endlabel);
       Label (elselabel)] @
      (trav fexp) @
      [InlineComment (Label (endlabel), node2str node)]

    (* Arrays *)
    | Allocate (dec, dims, _) ->
      let store = match (depthof dec, depthof node) with
      | (0, _)            -> Store (typeof dec, Glob,    indexof dec)
      | (a, b) when a = b -> Store (typeof dec, Current, indexof dec)
      | _                 -> raise InvalidNode
      in
      trav_all dims @
      [NewArray (basetypeof dec, List.length dims);
       InlineComment (store, node2str node)]

    | VarUse (dec, Some dims, _) ->
      let load = match (depthof dec, depthof node) with
      | (0, _)            -> Load (typeof dec, Glob,        indexof dec)
      | (a, b) when a = b -> Load (typeof dec, Current,     indexof dec)
      | (a, b)            -> Load (typeof dec, Rel (b - a), indexof dec)
      in
      (trav_all (List.rev dims)) @          (* push dimensions *)
      [InlineComment (load, nameof dec)] @  (* push array reference *)
      [InlineComment (LoadArray (basetypeof dec), node2str node)]

    | VarLet (dec, Some dims, value, _) ->
      let load = match (depthof dec, depthof node) with
      | (0, _)            -> Load (typeof dec, Glob,        indexof dec)
      | (a, b) when a = b -> Load (typeof dec, Current,     indexof dec)
      | (a, b)            -> Load (typeof dec, Rel (b - a), indexof dec)
      in
      (trav value) @                        (* push value *)
      (trav_all dims) @                     (* push dimensions *)
      [InlineComment (load, nameof dec)] @  (* push array reference *)
      [InlineComment (StoreArray (basetypeof dec), node2str node)]

    | _ -> raise InvalidNode
  in
  let instrs = trav program in

  (* Sort aggregated constants and add definitions
   * If possible, this should be rewritten in the future because it's a little
   * cumbersome right now... *)
  let pairs = ref [] in
  let add_pair value index =
    let com = sprintf "index %d" index in
    pairs := (InlineComment (ConstDef value, com), index) :: !pairs
  in
  Hashtbl.iter add_pair consts;
  let cmp (_, i) (_, j) = compare i j in
  let sorted_pairs = List.sort cmp !pairs in
  let const_defs = List.map (fun (d, _) -> d) sorted_pairs in

  instrs @ const_defs

let phase = function
  | Ast node -> Assembly (assemble node)
  | _ -> raise InvalidInput