open Printf
open Ast
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 (decl, dec_depth, _) ->
        (decl, 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 rec analyse scope depth args node =
    (* add_to_scope uses args, so it needs to be defined here *)
    let add_to_scope name decl depth scope =
        let (vars, funs) = scope in
        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 decl) msg args.verbose;
            prerr_loc_msg (locof orig) "Previously declared here:" args.verbose;
            raise EmptyError
        | Some _ ->
            Hashtbl.replace tbl name (decl, depth, name_type)
        | None ->
            Hashtbl.add tbl name (decl, depth, name_type)
    in

    let rec collect node = match node with
        (* Add node reference for this varname to vars map *)
        | VarDec (ctype, name, init, loc) ->
            let node = match init with
                | Some value -> VarDec (ctype, name, Some (collect value), loc)
                | None -> node
            in
            add_to_scope (Varname name) node depth scope;
            node

        (* For global vars, only add the name *)
        | GlobalDec (_, name, _)
        | GlobalDef (_, _, name, _, _) ->
            add_to_scope (Varname name) node depth scope;
            node

        (* Functions are traversed later on, for now only add the name *)
        | FunDec (_, name, _, _)
        | FunDef (_, _, name, _, _, _) ->
            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, _) ->
            let (decl, dec_depth) = check_in_scope (Varname name) node scope in
            VarUse (node, ctypeof decl, depth - dec_depth)

        | Deref (name, dims, loc) ->
            let (decl, dec_depth) = check_in_scope (Varname name) node scope in
            let node = Deref (name, List.map collect dims, loc) in
            VarUse (node, ctypeof decl, depth - dec_depth)

        | FunCall (name, args, loc) ->
            let (decl, dec_depth) = check_in_scope (Funcname name) node scope in
            let node = FunCall (name, transform_all collect args, loc) in
            FunUse (node, decl, depth - dec_depth)

        (* Assign statements are wrapped in VarLet nodes, which stores the type
         * and depth of the assigned variable are *)
        | Assign (name, None, value, loc) ->
            let (decl, dec_depth) = check_in_scope (Varname name) node scope in
            let assign = Assign (name, None, collect value, loc) in
            VarLet (assign, ctypeof decl, depth - dec_depth)

        | Assign (name, Some dims, value, loc) ->
            let (decl, dec_depth) = check_in_scope (Varname name) node scope in
            let dims = Some (List.map collect dims) in
            let assign = Assign (name, dims, collect value, loc) in
            VarLet (assign, ctypeof decl, depth - dec_depth)

        | _ -> transform_children collect node
    in

    (*let print_scope () =
        let (vars, funs) = scope in
        let print_key key value = prerr_string (" " ^ key) in
        prerr_string "vars: ";
        Hashtbl.iter print_key vars;
        prerr_endline "";
        prerr_string "funs: ";
        Hashtbl.iter print_key funs;
        prerr_endline "";
    in*)

    let rec traverse scope depth node =
        match node with
        (* Increase nesting level when entering function *)
        | FunDef (export, ret_type, name, params, body, loc) ->
            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) args body in
            FunDef (export, ret_type, name, params, body, loc)

        | Param (Array (_, dims), name, _) as node ->
            let rec add_dims = function
                | [] -> ()
                | Dim (name, _) as dim :: tail ->
                    add_to_scope (Varname name) (DimDec dim) depth scope;
                    add_dims tail
                | _ -> raise InvalidNode
            in
            add_dims dims;
            add_to_scope (Varname name) node 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;
     *)
    (*prerr_endline "";
    prerr_endline ("node:----\n" ^ Stringify.node2str node);
    prerr_endline "----";*)
    let node = collect node in
    (*prerr_endline "collected";
    print_scope ();
    prerr_endline "\ntraversing";*)

    let node = traverse scope depth node in
    (*prerr_endline "traversed";
    print_scope ();
    prerr_endline "";*)
    node

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

let rec phase input =
    prerr_endline "- Context analysis";
    match input with
    | Ast node -> Ast (analyse_context args node)
    | _ -> raise (InvalidInput "context analysis")