open Types
open Util
open Stringify

let tab = "    "
let max_instr_width = 26

let si = string_of_int

let op2str = function
  | Neg -> "neg"
  | Not -> "not"
  | Add -> "add"
  | Sub -> "sub"
  | Mul -> "mul"
  | Div -> "div"
  | Mod -> "rem"
  | Eq  -> "eq"
  | Ne  -> "ne"
  | Lt  -> "lt"
  | Le  -> "le"
  | Gt  -> "gt"
  | Ge  -> "ge"
  | _ -> raise (FatalError (Msg "operator unsupported by VM"))

let prefix = function
  | Bool _  -> "b"
  | Int _   -> "i"
  | Float _ -> "f"
  | Void    -> ""
  | _       -> "a"

let suffix = function
  | Glob        -> "g"
  | Current     -> ""
  | Local       -> "l"
  | Rel nesting -> "n " ^ si nesting

let rtn_suffix = function
  | ExternFun index    -> "e " ^ si index
  | LocalFun (size, label) -> " " ^ si size ^ " " ^ label

let rec instr2str = function
  (* Global / directives *)
  | Comment comment ->
    if Globals.args.verbose >= 2 then "; " ^ comment else ""
  | InlineComment (instr, comment) ->
    if Globals.args.verbose >= 2 then
      expand max_instr_width (instr2str instr) ^ " ; " ^ comment
    else
      instr2str instr
  | Label name ->
    name ^ ":"
  | Export (name, ret_type, arg_types, label) ->
    let types = List.map type2str (ret_type :: arg_types) in
    ".export \"" ^ name ^ "\" " ^ (String.concat " " types) ^ " " ^ label
  | Import (name, ret_type, arg_types) ->
    let types = List.map type2str (ret_type :: arg_types) in
    ".import \"" ^ name ^ "\" " ^ (String.concat " " types)
  | Global ctype ->
    ".global " ^ (type2str ctype)
  | ConstDef value ->
    ".const " ^ type2str (const_type value) ^ " " ^ const2str value

  (* Store *)
  | Store (ctype, scope, index) ->
    tab ^ prefix ctype ^ "store" ^ suffix scope ^ " " ^ si index

  (* Load *)
  | Load (ctype, Current, index) when index >= 0 & index <= 3 ->
    tab ^ prefix ctype ^ "load_" ^ si index
  | Load (ctype, scope, index) ->
    tab ^ prefix ctype ^ "load" ^ suffix scope ^ " " ^ si index
  | LoadConst (ctype, index) ->
    tab ^ prefix ctype ^ "loadc " ^ si index
  | LoadImm (BoolVal b) ->
    tab ^ "bloadc_" ^ (if b then "t" else "f")
  | LoadImm (IntVal i) when i < 0l ->
    tab ^ "iloadc_m" ^ Int32.to_string (Int32.neg i)
  | LoadImm (IntVal i) ->
    tab ^ "iloadc_" ^ Int32.to_string i
  | LoadImm (FloatVal i) ->
    tab ^ "floadc_" ^ si (int_of_float i)

  (* Operators *)
  | Op (op, ctype) ->
    tab ^ prefix ctype ^ op2str op
  | Convert (src, tgt) ->
    tab ^ prefix src ^ "2" ^ prefix tgt
  | Inc (index, const) ->
    tab ^ "iinc " ^ si index ^ " " ^ si const
  | Dec (index, const) ->
    tab ^ "idec " ^ si index ^ " " ^ si const
  | IncOne index ->
    tab ^ "iinc_1 " ^ si index
  | DecOne index ->
    tab ^ "idec_1 " ^ si index

  (* Control flow *)
  | RtnInit scope ->
    tab ^ "isr" ^ suffix scope
  | RtnJmp scope ->
    tab ^ "jsr" ^ rtn_suffix scope
  | RtnEnter stack_len ->
    tab ^ "esr " ^ si stack_len
  | Ret ctype ->
    tab ^ prefix ctype ^ "return"
  | Branch (true, target) ->
    tab ^ "branch_t " ^ target
  | Branch (false, target) ->
    tab ^ "branch_f " ^ target
  | Jump target ->
    tab ^ "jump " ^ target

  (* Stack management *)
  | Pop ctype ->
    tab ^ prefix ctype ^ "pop"

  (* Arrays *)
  | NewArray (basetype, ndims) ->
    tab ^ prefix basetype ^ "newa " ^ si ndims
  | ArraySize index ->
    tab ^ "asize " ^ si index
  | LoadArray basetype ->
    tab ^ prefix basetype ^ "loada"
  | StoreArray basetype ->
    tab ^ prefix basetype ^ "storea"

  | EmptyLine -> ""

(* Print assembly instructions to the given file pointer. *)
let rec print_assembly oc instrs =
  let output_line line =
    output_string oc line;
    output_char oc '\n';
  in
  let endbuf = ref [] in
  let rec trav = function
    | [] -> ()
    | EmptyLine :: tl -> output_line ""; trav tl
    | hd :: tl ->
      let line = instr2str hd in
      begin
        if String.length line > 0 && line.[0] = '.' then
          endbuf := line :: !endbuf
        else if String.length line > 0 then
          output_line line
      end;
      trav tl
  in
  trav instrs;
  if List.length !endbuf > 1 then begin
    output_line (instr2str (Comment ("globals:")));
    let cmp a b = compare (String.sub b 0 7) (String.sub a 0 7) in
    List.iter output_line (List.sort cmp (List.rev !endbuf))
  end

let phase = function
  | Ast node as input ->
    prerr_endline hline;
    prerr_endline (node2str node);
    prerr_endline hline;
    input

  | FileContent (display_name, content) as input ->
    prerr_endline hline;
    prerr_endline (display_name ^ ":\n");
    prerr_endline content;
    prerr_endline hline;
    input

  | Assembly instrs as input ->
    prerr_endline hline;
    print_assembly stderr instrs;
    prerr_endline hline;
    input

  | Empty -> Empty