open Types
open Util
open Globals

let create_param ctype name =
    let param = Param (ctype, name, [Depth 1]) in
    let value = VarUse (param, None, [Type ctype; Depth 1]) in
    (param, value)

let call node args depth = match node with
    | FunDec (ctype, name, _, _) as dec ->
        FunUse (dec, args, [Type ctype; Depth depth])
    | _ -> raise InvalidNode

let process globals = function
    | GlobalDef (true, Array ctype, name, None, ann) as dec ->
        (* 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 = VarUse (dec, Some [index], [Type ctype; Depth 1]) in
        let body = Block [Return (var, [])] in
        let getter = FunDef (true, ctype, name ^ "$get", [param], body, []) 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 (dec, Some [index], value, [])] in
        let setter = FunDef (true, Void, name ^ "$set", [param1; param2], body, []) in

        [getter; setter]

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

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

        [getter; setter]

    | GlobalDec (Array ctype, name, ann) ->
        (* Getters for external array variable: create getter and setter for a
         * given index *)
        let (param, _) = create_param Int "index" in
        let getter = FunDec (ctype, name ^ "$get", [param], []) 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], []) in

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

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

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

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

    | _ -> raise InvalidNode

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

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

    | GlobalDec (ctype, name, ann) 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, ann) ->
        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, ann)

    | VarUse (dec, None, _) when Hashtbl.mem scope (nameof dec) ->
        let (get, _) = Hashtbl.find scope (nameof dec) in
        get [] depth

    | VarLet (dec, dims, value, _) when Hashtbl.mem scope (nameof dec) ->
        let dims = optmapl (replace_vars scope depth) dims in
        let (_, set) = Hashtbl.find scope (nameof dec) in
        Expr (set (dims @ [replace_vars scope depth value]) depth)

    | node -> transform_children (replace_vars scope depth) node

let phase = function
    | Ast node ->
        let globals = Hashtbl.create 20 in
        let node = create_funcs globals node in
        Ast (replace_vars globals 0 node)
    | _ -> raise (InvalidInput "extern vars")