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_args callstack node =
        let trav = trav_args callstack in
        let rec trav_all = function
            | [] -> []
            | hd :: tl -> trav hd @ (trav_all tl)
        in
        let rec traverse_localfuns callstack = function
            | LocalFuns body -> List.concat (List.map (trav_args callstack) body)
            | Block body -> List.concat (List.map (traverse_localfuns callstack) body)
            | _ -> []
        in
        match node with
        (* Global *)
        | Program (decls, _) ->
            trav_all decls

        | GlobalDef (_, ctype, _, _, _) ->
            [Global ctype]

        | FunDec (ret_type, name, params, _) ->
            [Import (name, ret_type, List.map typeof params)]

        | FunDef (export, ret_type, name, params, body, _) ->
            let callstack = name :: callstack in
            let label = String.concat "$" (List.rev callstack) in
            (if export then
                let param_types = List.map typeof params in
                [Export (name, ret_type, param_types, label)]
            else []) @
            [
                Comment ("function \"" ^ label ^ "\":");
                Label label;
                RtnEnter (indexof node);
            ] @
            (trav_args callstack body) @
            (match ret_type with Void -> [Ret Void] | _ -> []) @
            [EmptyLine] @
            (traverse_localfuns callstack body)

        | VarDec (ctype, name, _, _) ->
            [comline (sprintf "index %d: %s %s" (indexof node) (type2str ctype) name)]

        | LocalFuns _ -> []

        | Block body | VarDecs body -> trav_all body

        (* Statements *)
        | VarLet (dec, None, value, _) ->
            let store = match (depthof dec, depthof node) with
                | (0, _)            -> StoreGlob (typeof dec, indexof dec)
                | (a, b) when a = b -> StoreLoc  (typeof dec, indexof dec)
                | (a, b)            -> StoreRel  (typeof dec, b - a, indexof dec)
            in
            trav value @ [InlineComment (store, node2str node)]

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

        (* Expressions *)
        | Const (BoolVal _, _) ->
            [LoadImm node]

        | Const (value, _) ->
            let def = if Hashtbl.mem consts value then [] else (
                Hashtbl.add consts value true;
                [ConstDef value]
            ) in
            def @ [LoadConst (typeof node, indexof node)]

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

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

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

        | TypeCast (ctype, value, _) ->
            let vtype = typeof value in
            (match (ctype, vtype) with
            | (Float, Int) | (Int, Float) -> ()
            | _ -> raise (NodeError (node, "invalid typecast"))
            );
            trav value @ [Convert (vtype, ctype)]

        | _ -> [Comment ("FIXME: " ^ Stringify.node2str node)]
        (*| _ -> raise InvalidNode*)
    in
    trav_args [] program

let rec phase input =
    log_line 1 "- Assembly";
    match input with
    | Ast node -> Assembly (assemble node)
    | _ -> raise (InvalidInput "assembly")