open Types
open Util
open Stringify

let tab = "    "

let si = string_of_int

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)

let ctype2str = Stringify.type2str
let type2str = function
    | Array (t, dims) -> ctype2str t ^ repeat "," (List.length dims - 1)
    | t -> ctype2str t

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 (CompileError ("operator unsupported by VM"))

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

let rec instr2str = function
    (* Global / directives *)
    | Comment comment ->
        if args.verbose >= 2 then "# " ^ comment else ""
    | InlineComment (instr, comment) ->
        if args.verbose >= 2 then
            expand 20 (instr2str instr) ^ "# " ^ comment
        else ""
    | 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 node ->
        ".const " ^ (const2str node)

    (* Store *)
    | StoreGlob (ctype, index) ->
        tab ^ prefix ctype ^ "storeg " ^ si index
    | StoreLoc (ctype, index) ->
        tab ^ prefix ctype ^ "store " ^ si index
    | StoreRel (ctype, nesting, index) ->
        tab ^ prefix ctype ^ "storen " ^ si nesting ^ " " ^ si index

    (* Load *)
    | LoadGlob (ctype, index) ->
        tab ^ prefix ctype ^ "loadg " ^ si index
    | LoadLoc (ctype, index) ->
        tab ^ prefix ctype ^ "load " ^ si index
    | LoadRel (ctype, nesting, index) ->
        tab ^ prefix ctype ^ "loadn " ^ si nesting ^ " " ^ si index

    | LoadConst (ctype, index) ->
        tab ^ prefix ctype ^ "loadc " ^ si index
    | LoadImm (Const (BoolVal b, _)) ->
        tab ^ "bloadc_" ^ (if b then "t" else "f")
    | LoadImm (Const (IntVal i, _)) when i < 0 ->
        tab ^ "iloadc_m" ^ si (-i)
    | LoadImm (Const (IntVal i, _)) ->
        tab ^ "iloadc_" ^ si i
    | LoadImm (Const (FloatVal i, _)) ->
        tab ^ "floadc_" ^ si (int_of_float i)

    | Op (op, ctype) ->
        tab ^ prefix ctype ^ op2str op

    | Convert (src, tgt) ->
        tab ^ prefix src ^ "2" ^ prefix tgt

    (* Control flow *)
    | RtnEnter stack_len ->
        tab ^ "esr " ^ si stack_len
    | Ret ctype ->
        tab ^ prefix ctype ^ "return"

    | EmptyLine -> ""
    | DummyInstr -> tab ^ "<dummy>"
    | _ -> tab ^ "<unknown instruction>"

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
        | [] -> ()
        | hd :: tl ->
            let line = instr2str hd in
            (if String.length line > 0 && line.[0] = '.' then
                endbuf := line :: !endbuf
            else
                output_line line
            );
            trav tl
    in
    trav instrs;
    if List.length !endbuf > 1 then (
        output_line (instr2str (Comment ("globals:")));
        let cmp a b = compare (String.sub a 0 7) (String.sub b 0 7) in
        List.iter output_line (List.sort cmp (List.rev !endbuf))
    ); ()

let phase = function
    | Ast node as input ->
        if args.verbose >= 2 then (
            prerr_endline "--------------------------------------------------";
            prerr_endline (node2str node);
            prerr_endline "--------------------------------------------------"
        );
        input

    | FileContent (display_name, content) as input ->
        if args.verbose >= 2 then (
            prerr_endline "--------------------------------------------------";
            prerr_endline (display_name ^ ":\n");
            prerr_endline content;
            prerr_endline "--------------------------------------------------"
        );
        input

    | Assembly instrs as input ->
        (match args.outfile with
        | Some filename ->
            let oc = open_out filename in
            print_assembly oc instrs;
            close_out oc
        | None ->
            if args.verbose >= 2 then
                prerr_endline "--------------------------------------------------";

            print_assembly stdout instrs;
        );
        input

    | _ -> raise (InvalidInput "print")