open Printf
open Types
open Util

type nametype = Varname of string | Funcname of string

let type2str = function Funcname _ -> "function" | Varname _ -> "variable"

let mapfind name tbl =
    if Hashtbl.mem tbl name then Some (Hashtbl.find tbl name) else None

let check_in_scope name errnode scope =
    let (vars, funs) = scope in
    let (name, tbl, other_map, desired_type) = match name with
        | Varname  name -> (name, vars, funs, "variable")
        | Funcname name -> (name, funs, vars, "function")
    in
    match mapfind name tbl with
    | Some (dec, dec_depth, _) ->
        (dec, dec_depth)
    | None ->
        let msg = match mapfind name other_map with
            | Some _ -> sprintf "\"%s\" is not a %s" name desired_type
            | None   -> sprintf "undefined %s \"%s\"" desired_type name
        in
        raise (NodeError (errnode, msg))

let add_to_scope name dec depth (vars, funs) =
    prerr_string "__add_to_scope: "; prt_node dec;
    let (name, tbl, name_type) = match name with
        | Varname name  -> (name, vars, "variable")
        | Funcname name -> (name, funs, "function")
    in
    match mapfind name tbl with
    (* Identifiers of lower depth may be overwritten, but idenetifiers at
        * the same depth must be unique for consistency *)
    | Some (orig, orig_depth, _) when orig_depth >= depth ->
        let msg = sprintf "Error: cannot redeclare %s \"%s\"" name_type name in
        prerr_loc_msg (locof dec) msg;
        prerr_loc_msg (locof orig) "Previously declared here:";
        raise EmptyError
    | Some _ ->
        Hashtbl.replace tbl name (dec, depth, name_type)
    | None ->
        Hashtbl.add tbl name (dec, depth, name_type)

let rec analyse scope depth node =
    let rec collect node = match node with
        (* For variables, add the name (array dimensions are added
         * implicitly, since they have separate VarDec nodes which were added
         * during the desugaring phase *)
        | VarDec (_, name, _, _)
        | GlobalDec (_, name, _)
        | GlobalDef (_, _, name, _, _) ->
            let node = annotate (Depth depth) node in
            add_to_scope (Varname name) node depth scope;
            node

        (* Functions are traversed later on, for now only add the name *)
        | FunDec (_, name, _, _)
        | FunDef (_, _, name, _, _, _) ->
            let node = annotate (Depth depth) node in
            add_to_scope (Funcname name) node depth scope;
            node

        (* For a variable or function call, look for its declaration in the
         * current scope and save a its type/depth information  *)
        | Var (name, dims, ann) ->
            let (dec, dec_depth) = check_in_scope (Varname name) node scope in
            VarUse (dec, optmap collect dims, Depth depth :: ann)

        | FunCall (name, args, ann) ->
            let (dec, dec_depth) = check_in_scope (Funcname name) node scope in
            FunUse (dec, List.map collect args, Depth depth :: ann)

        (* Assign statements are replaced with VarLet nodes, which stores the
         * declaration of the assigned variable *)
        | Assign (name, dims, value, ann) ->
            let (dec, dec_depth) = check_in_scope (Varname name) node scope in
            VarLet (dec, optmap collect dims, collect value, Depth depth :: ann)

        | Allocate (dec, dims, ann) ->
            let (dec, dec_depth) = check_in_scope (Varname (nameof dec)) node scope in
            Allocate (dec, List.map collect dims, Depth depth :: ann)

        | _ -> transform_children collect node
    in

    let rec traverse scope depth node =
        match node with
        (* Increase nesting level when entering function *)
        | FunDef (export, ret_type, name, params, body, ann) ->
            let (vars, funs) = scope in
            let local_scope = (Hashtbl.copy vars, Hashtbl.copy funs) in
            let params = List.map (traverse local_scope (depth + 1)) params in
            let body = analyse local_scope (depth + 1) body in
            FunDef (export, ret_type, name, params, body, ann)

        | Param (ArrayDims (ctype, dims), name, ann) ->
            let rec add_dims = function
                | [] -> []
                | Dim (name, ann) :: tl ->
                    let dim = Dim (name, Depth depth :: ann) in
                    add_to_scope (Varname name) dim depth scope;
                    dim :: (add_dims tl)
                | _ -> raise InvalidNode
            in
            let node = Param (ArrayDims (ctype, add_dims dims), name, ann) in
            add_to_scope (Varname name) node depth scope;
            node

        | VarDec _ -> node

        | Param (_, name, _) ->
            let node = annotate (Depth depth) node in
            add_to_scope (Varname name) node depth scope;
            node

        (* Do not traverse into external function declarations, since their
         * parameters must not be added to the namespace *)
        | FunDec _ -> node

        | _ -> transform_children (traverse scope depth) node
    in

    (*
     * First collect all definitions at the current depth. Then, traverse into
     * functions with a copy of the current scope. This is needed because
     * functions can access all identifiers in their surrounding scope.
     * E.g., the following is allowed:
     *
     * void foo() { glob = 1; }
     * int glob;
     *)
    let node = collect node in

    let node = traverse scope depth node in
    node

let analyse_context program =
    let scope = (Hashtbl.create 20, Hashtbl.create 20) in
    analyse scope 0 program

let phase = function
    | Ast node -> Ast (analyse_context node)
    | _ -> raise (InvalidInput "context analysis")