open Printf
open Ast
open Util

module StrMap = Map.Make (String)

let analyse_context args node =
    let scope = ref StrMap.empty in
    let add_to_scope name decl depth desc =
        if StrMap.mem name !scope then (
            let msg = sprintf "Error: cannot redeclare %s \"%s\"" desc name in
            prerr_loc_msg (locof !decl) msg args.verbose;

            let (orig, _) = StrMap.find name !scope in
            prerr_loc_msg (locof !orig) "Previously declared here:" args.verbose;

            raise EmptyError
        ) else
            scope := StrMap.add name (decl, depth) !scope
    in
    let rec analyse depth 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 ->
                    let value = analyse depth value in
                    VarDec (ctype, name, Some value, loc)
                | None -> node
            in
            add_to_scope name (ref node) depth "variable";
            node

        (* For a variable, look for its declaration in the current scope and
         * save a reference with the relative nesting depth *)
        | Var (name, _) ->
            if StrMap.mem name !scope then
                let (decl, decl_depth) = StrMap.find name !scope in
                VarUse (node, decl, depth - decl_depth)
            else
                raise (NodeError (node, (sprintf "undefined variable \"%s\"" name)))

        (* Increase nesting level when entering function *)
        | FunDef (export, ret_type, name, params, body, loc) ->
            add_to_scope name (ref node) depth "function";
            let params = List.map (analyse (depth + 1)) params in
            let body = analyse (depth + 1) body in
            FunDef (export, ret_type, name, params, body, loc)

        | Param (ArrayDec (_, dims) as atype, name, _) as node ->
            let add dim = add_to_scope dim (ref (Type atype)) depth "variable" in
            List.iter add dims;
            add_to_scope name (ref node) depth "variable";
            node

        | Param (_, name, _) ->
            add_to_scope name (ref node) depth "variable";
            node

        | node -> transform_children (analyse depth) node
    in
    analyse 0 node

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