open Ast
open Util

(* Remove Type nodes temporarily for easy traversal *)
let rec prune_types = function
    | Type (node, _) -> prune_types node
    | node -> transform_children prune_types node

let create_param ctype name =
    let param = Param (ctype, name, noloc) in
    let value = VarUse (Var (name, noloc), ctype, 0) in
    (param, value)

let call node args depth = match node with
    | FunDec (ctype, name, _, _) as def ->
        Type (FunUse (FunCall (name, args, noloc), def, depth), ctype)
    | _ -> raise InvalidNode

let process globals = function
    | GlobalDef (true, Array (ctype, dims), name, None, loc) ->
        (* Getters for array variable: crate getter for given index Note that
         * getters and setters for dimensions are automatically generated,
         * because they have been put into new global variables during the
         * desugarin phase *)
        let (param, index) = create_param Int (fresh_var "index") in
        let var = Deref (name, [index], noloc) in
        let body = Block [Return (var, noloc)] in
        let getter = FunDef (true, ctype, name ^ "$get", [param], body, noloc) in

        (* Setters for array variable: create setter for given index *)
        let (param1, index) = create_param Int (fresh_var "index") in
        let (param2, value) = create_param ctype (fresh_var "value") in
        let body = Block [VarLet (Assign (name, Some [index], value, noloc), ctype, 1)] in
        let setter = FunDef (true, Void, name ^ "$set", [param1; param2], body, noloc) in

        [getter; setter]

    | GlobalDef (true, ctype, name, None, loc) ->
        (* Getter for basic variable type: return the variable *)
        let var = VarUse (Var (name, noloc), ctype, 1) in
        let body = [Return (var, noloc)] in
        let getter = FunDef (true, ctype, name ^ "$get", [], Block body, noloc) in

        (* Setter for basic variable type: assign the variable *)
        let (param, value) = create_param ctype (fresh_var "value") in
        let body = [VarLet (Assign (name, None, value, noloc), ctype, 1)] in
        let setter = FunDef (true, Void, name ^ "$set", [param], Block body, noloc) in

        [getter; setter]

    | GlobalDec (Array (ctype, dims), name, loc) ->
        (* Getters for external array variable: create getter for a given index *)
        (* Setters for external array variable:
         * - define setter for a given index
         * - define setter for each dimension*)
        let (param, _) = create_param Int "index" in
        let rec process_dims i = function
            | [] -> []
            | Dim (oldname, _) :: tl ->
                let dimname = name ^ "$" ^ string_of_int i in

                let getter = FunDec (Void, dimname ^ "$get", [], noloc) in

                let (param, _) = create_param ctype "value" in
                let setter = FunDec (Void, dimname ^ "$set", [param], noloc) in

                Hashtbl.add globals oldname (call getter, call setter);
                getter :: setter :: (process_dims (i + 1) tl)
            | _ -> raise InvalidNode
        in

        let getter = FunDec (ctype, name ^ "$get", [param], noloc) in

        let (param1, index) = create_param Int "index" in
        let (param2, value) = create_param ctype "value" in
        let setter = FunDec (Void, name ^ "$set", [param1; param2], noloc) in

        Hashtbl.add globals name (call getter, call setter);
        getter :: setter :: (process_dims 1 dims)

    (* Getter for basic variable type: return the variable *)
    | GlobalDec (ctype, name, loc) ->
        let getter = FunDec (ctype, name ^ "$get", [], noloc) in

        let (param, _) = create_param ctype "value" in
        let setter = FunDec (Void, name ^ "$set", [param], noloc) in

        Hashtbl.add globals name (call getter, call setter);
        [getter; setter]

    | _ -> raise InvalidNode

let rec create_funcs globals = function
    | Program (decls, loc) ->
        let decls = List.map (create_funcs globals) decls in
        Program (flatten_blocks (List.map (create_funcs globals) decls), loc)

    | GlobalDef (true, ctype, name, None, loc) as node ->
        Block (GlobalDef (false, ctype, name, None, loc) ::
               (process globals node))

    | GlobalDec (ctype, name, loc) as node ->
        Block (process globals node)

    | node -> transform_children (create_funcs globals) node

let rec replace_vars scope depth = function
    | (VarDec (_, name, _, _) as node)
    | (Param (_, name, _) as node)
            when Hashtbl.mem scope name ->
        Hashtbl.remove scope name;
        node

    | FunDef (export, ret_type, name, params, body, loc) ->
        let local_scope = Hashtbl.copy scope in
        let trav = replace_vars local_scope (depth + 1) in
        let params = List.map trav params in
        FunDef (export, ret_type, name, params, trav body, loc)

    | VarUse (Var (name, loc), _, _) as node when Hashtbl.mem scope name ->
        let (get, _) = Hashtbl.find scope name in
        get [] depth

    | VarLet (Assign (name, None, value, _), _, _) when Hashtbl.mem scope name ->
        let (_, set) = Hashtbl.find scope name in
        Expr (set [replace_vars scope depth value] depth)

    | VarLet (Assign (name, Some dims, value, _), _, _) when Hashtbl.mem scope name ->
        let dims = List.map (replace_vars scope depth) dims in
        let (_, set) = Hashtbl.find scope name in
        Expr (set (dims @ [replace_vars scope depth value]) depth)

    | node -> transform_children (replace_vars scope depth) node

let rec phase input =
    prerr_endline "- Create getters and setters for extern variables";
    match input with
    | Ast node ->
        let globals = Hashtbl.create 20 in
        let node = create_funcs globals (prune_types node) in
        Ast (replace_vars globals 0 node)
    | _ -> raise (InvalidInput "extern vars")