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 pad width s = s ^ (repeat " " (String.length s - width))
let paddall width = List.map (pad width)

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

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

let instr2str = function
    (* Global / directives *)
    | Comment comment ->
        "# " ^ comment
    | 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)
    | ConstDef node ->
        ".const  " ^ (const2str node)
    | Global ctype ->
        ".global " ^ (type2str ctype)

    (* 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 constant *)
    | 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)

    (* 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 8) (String.sub b 0 8) 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")