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) =
    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
        (* Add node reference for this varname to vars map *)
        | VarDec (ctype, name, init, ann) ->
            let node = match init with
                | Some value -> VarDec (ctype, name, Some (collect value),
                                        Depth depth :: ann)
                | None -> VarDec (ctype, name, init, Depth depth :: ann)
            in
            add_to_scope (Varname name) node depth scope;
            node

        (* For global vars, only the name and array dimensions *)
        | GlobalDec (ctype, name, ann) ->
            let ctype = match ctype with
                | Array (ctype, dims) -> Array (ctype, List.map collect dims)
                | _ -> ctype
            in
            let node = GlobalDec (ctype, name, Depth depth :: ann) in
            add_to_scope (Varname name) node depth scope;
            node

        | Dim (name, ann) ->
            let node = Dim (name, Depth depth :: ann) in
            add_to_scope (Varname name) node depth scope;
            node

        | GlobalDef (export, ctype, name, init, ann) ->
            let ctype = match ctype with
                | Array (ctype, dims) -> Array (ctype, List.map collect dims)
                | _ -> ctype
            in
            let node = GlobalDef (export, ctype, name, init, Depth depth :: ann) in
            add_to_scope (Varname name) node depth scope;
            node

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

        | FunDef (export, ret_type, name, params, body, ann) ->
            let node = FunDef (export, ret_type, name, params, body, Depth depth :: ann) 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 type
         * and depth of the assigned variable are *)
        | 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 (Array (ctype, dims), name, ann) as node ->
            let _ = List.map (traverse scope depth) dims in
            add_to_scope (Varname name) node depth scope;
            node

        | Dim (name, _) as dim ->
            add_to_scope (Varname name) dim depth scope;
            node

        | Param (_, name, _) ->
            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 rec phase input =
    log_line 1 "- Context analysis";
    match input with
    | Ast node -> Ast (analyse_context node)
    | _ -> raise (InvalidInput "context analysis")