open Xi_lib
open Ast
(* W Xi_lib.Types są definicje typów i środowiska typowego *)
open Types

module Make() = struct

  (* Logger: używa się go jak Format.printf *)
  let logf fmt = Logger.make_logf __MODULE__ fmt

  module Check () = struct

    (* Zgłaszaczka błędów *)
    module ErrorReporter = Typechecker_errors.MakeOneShotErrorReporter ()

    (* Hashtablica którą zwracamy jak wszystko jest OK.
     * Mapuje znacznik węzła na przypisany typ. Dzięki tej tablicy
     * późniejsze etapy kompilatora będą miały dostęp do policzonych
     * typów wyrażeń 
     * 
     * Jeżeli typowanie się powiedzie to zawartość tablicy wydrukuje się
     * do pliku xilog/004.typechecking.types
     *)
    let node2type_map = Hashtbl.create 513

    let rec te2type_ret = function
      | TEXPR_Int _ -> TP_Int
      | TEXPR_Bool _ -> TP_Bool
      | TEXPR_Array {sub;dim=None;_} -> TP_Array (te2type_ret sub)
      | TEXPR_Array {sub;dim;loc} -> ErrorReporter.report_array_initialization_forbidden ~loc
    (* --------------------------------------------------- *)
    (* Funkcja nakładka na inferencję, jej zadanie to uzupełniać hashtablicę node2type_map *)
    let rec infer_expression env e =
      let tp = _infer_expression env e in
      Hashtbl.replace node2type_map (tag_of_expression e) tp;
      logf "%s: inferred type %s"
        (string_of_location
        (location_of_expression e))
        (string_of_normal_type tp);
      tp

    (* --------------------------------------------------- *)
    (* Oddolna strategia *)
    and _infer_expression env = function
      | EXPR_Id {id; loc; _} ->
        begin
          match TypingEnvironment.lookup id env with
          | None -> ErrorReporter.report_unknown_identifier ~loc ~id
          | Some tp -> match tp with
                       | ENVTP_Var tp -> tp
                       | ENVTP_Fn (_,_) -> ErrorReporter.report_identifier_is_not_variable ~loc ~id
        end

      | EXPR_Int _ ->
        TP_Int

      | EXPR_Char _ ->
        TP_Int

      | EXPR_Bool _ ->
        TP_Bool

      | EXPR_Index {expr;index;loc; _} ->
        check_expression env  TP_Int index;
        begin
          match  infer_expression env expr with
          | TP_Array tp -> tp
          | TP_Int   -> ErrorReporter.report_expected_array ~loc ~actual:TP_Int
          | TP_Bool  -> ErrorReporter.report_expected_array ~loc ~actual:TP_Bool
        end
      | EXPR_Call call ->
        check_function_call env call

      | EXPR_Length {arg;loc;_} ->
        begin
          match infer_expression env arg with
          | TP_Array _ -> TP_Int
          | TP_Int   -> ErrorReporter.report_expected_array ~loc ~actual:TP_Int
          | TP_Bool  -> ErrorReporter.report_expected_array ~loc ~actual:TP_Bool
        end

      | EXPR_Relation {lhs; rhs; op=RELOP_Ge; _} 
      | EXPR_Relation {lhs; rhs; op=RELOP_Gt; _} 
      | EXPR_Relation {lhs; rhs; op=RELOP_Lt; _} 
      | EXPR_Relation {lhs; rhs; op=RELOP_Le; _}  ->
        check_expression env TP_Int lhs;
        check_expression env TP_Int rhs;
        TP_Bool

      | EXPR_Relation {lhs; rhs; op=RELOP_Eq; _} 
      | EXPR_Relation {lhs; rhs; op=RELOP_Ne; _} ->
        let tp = infer_expression env lhs in
        check_expression env tp rhs;
        TP_Bool
        (* Consider backtracking *)

        (* Reguła dla dodawania, jak w treści zadania *)
      | EXPR_Binop {loc; lhs; rhs; op=BINOP_Add; _} ->
        begin match infer_expression env lhs with
        | (TP_Array _) as tp
        | (TP_Int as tp) ->
          check_expression env tp rhs;
          tp
        | _ ->
          let descr = "operator + expects integer or array" in
         ErrorReporter.report_other_error ~loc ~descr
        end

      | EXPR_Binop {lhs; rhs; op=BINOP_And;_} 
      | EXPR_Binop {lhs; rhs; op=BINOP_Or; _} ->
        check_expression env TP_Bool lhs;
        check_expression env TP_Bool rhs;
        TP_Bool

      | EXPR_Binop {lhs; rhs; op=BINOP_Sub;_} 
      | EXPR_Binop {lhs; rhs; op=BINOP_Rem;_} 
      | EXPR_Binop {lhs; rhs; op=BINOP_Mult;_} 
      | EXPR_Binop {lhs; rhs; op=BINOP_Div; _} ->
        check_expression env TP_Int lhs;
        check_expression env TP_Int rhs;
        TP_Int

      | EXPR_Unop {op=UNOP_Neg; sub; _} ->
        check_expression env TP_Int sub;
        TP_Int

      | EXPR_Unop {op=UNOP_Not; sub; _} ->
        check_expression env TP_Bool sub;
        TP_Bool

      | EXPR_String _ ->
        TP_Array TP_Int

      | EXPR_Struct {elements=[]; loc; _} ->
        ErrorReporter.report_cannot_infer ~loc

      | EXPR_Struct {elements=x::xs; _} ->
        let t=infer_expression env x in
        let _=List.map (check_expression env t) xs in
        TP_Array t

    and check_function_call env (Call {loc;callee;_} as call)=
      match check_function_args env call with
      | [] -> ErrorReporter.report_other_error ~loc ~descr:("Procedures don't return expressions")
      | [x] ->x
      | _  -> ErrorReporter.report_expected_function_returning_one_value ~loc ~id:callee
    and check_function_args env (Call {loc;callee;arguments;_})  =
      begin
      match TypingEnvironment.lookup callee env with
        | None -> ErrorReporter.report_unknown_identifier ~loc ~id:callee
        | Some ENVTP_Var _ -> ErrorReporter.report_identifier_is_not_callable ~loc ~id:callee
        | Some (ENVTP_Fn (arg_types,ret_types)) ->
        begin
          let rec aux args types = match args,types with
            | [],[] -> ()
            | arg::args, t::ts -> check_expression env t arg;aux args ts
            | _,[] -> ErrorReporter.report_bad_number_of_arguments ~loc ~expected:(List.length arg_types) ~actual:(List.length arguments)
            | [],_ -> ErrorReporter.report_bad_number_of_arguments ~loc ~expected:(List.length arg_types) ~actual:(List.length arguments)
          in aux arguments arg_types;
          ret_types
        end
      end
    and te2base env = begin function
        | TEXPR_Int _ -> TP_Int
        | TEXPR_Bool _ -> TP_Bool
        | TEXPR_Array {sub;dim=None;loc} -> ErrorReporter.report_array_initialization_forbidden ~loc
        | TEXPR_Array {sub;dim=Some x;_} -> check_expression env TP_Int x;TP_Array (te2base env sub)
      end
    and te2type_decl env = begin function
        | TEXPR_Int _ -> TP_Int
        | TEXPR_Bool _ -> TP_Bool
        | TEXPR_Array {sub;dim=None;_} -> TP_Array (te2type_decl env sub)
        | TEXPR_Array {sub;dim=Some x;_} -> check_expression env TP_Int x;TP_Array (te2base env sub)
      end
    and add_var_decl env (VarDecl {loc;id;tp}) = 
        match TypingEnvironment.add id (ENVTP_Var (te2type_decl env tp)) env with
        | (env,true) -> env
        | (env,false) -> ErrorReporter.report_shadows_previous_definition ~loc ~id
    (* --------------------------------------------------- *)
    (* Odgórna strategia: zapish w node2type_map oczekiwanie a następnie
     * sprawdź czy nie zachodzi specjalny przypadek. *)
    and check_expression env expected e =
      logf "%s: checking expression against %s"
        (string_of_location (location_of_expression e))
        (string_of_normal_type expected);
      Hashtbl.replace node2type_map (tag_of_expression e) expected;

      (* Sprawdzamy specjalne przypadki *)
      match e, expected with
          (* Mamy sprawdzić `{elements...}` kontra `tp[]`, czyli sprawdzamy 
          * elementy kontra typ elementu tablicy `tp` *)
      | EXPR_Struct {elements; _}, TP_Array tp ->
        List.iter (check_expression env tp) elements
      | EXPR_Binop {op=BINOP_Add; lhs; rhs;_}, TP_Int ->
        check_expression env TP_Int lhs;
        check_expression env TP_Int rhs
      | EXPR_Binop {op=BINOP_Add; lhs; rhs;_}, (TP_Array tp as t) ->
        check_expression env t lhs;
        check_expression env t rhs
      | EXPR_Index {expr;index;loc; _}, tp ->
        check_expression env  TP_Int index;
        check_expression env (TP_Array tp) expr
      (* ========== !! Zaimplementuj pozostale przypadki !! =========  *)

      (* Fallback do strategii oddolnej *)
      | _ ->
        let actual = infer_expression env e in
        if actual <> expected then
          ErrorReporter.report_type_mismatch
            ~loc:(location_of_expression e)
            ~actual
            ~expected

    (* --------------------------------------------------- *)
    (* Pomocnicza funkcja do sprawdzania wywołania procedury *)

    let check_procedure_call env (Call {loc;_} as call) : unit = 
        match check_function_args env call with
        | [] -> ()
        | _ -> ErrorReporter.report_other_error ~loc ~descr:"Function used as procedure"

    (* --------------------------------------------------- *)
    (* Rekonstrukcja typu dla lvalue *)

    let infer_lvalue env = function
      | LVALUE_Id {id;loc;_} -> 
      begin
        match TypingEnvironment.lookup id env with
        | None -> ErrorReporter.report_unknown_identifier ~loc ~id
        | Some ENVTP_Fn _ -> ErrorReporter.report_identifier_is_not_variable ~loc ~id
        | Some ENVTP_Var tp ->tp
      end

      | LVALUE_Index {index; sub; loc} ->
        check_expression env TP_Int index;
        begin
          match infer_expression env sub with
            | TP_Int -> ErrorReporter.report_expected_array ~loc ~actual:TP_Int
            | TP_Bool -> ErrorReporter.report_expected_array ~loc ~actual:TP_Bool
            | TP_Array tp -> tp
        end


    (* --------------------------------------------------- *)
    (* Sprawdzanie statementów *)

    let rec check_statement env = function
      (* Proste, wyinferuj typ `lhs` i sprawdź `rhs` kontra wynik *)
      | STMT_Assign {lhs; rhs; _} ->
        let lhs_tp = infer_lvalue env lhs in
        check_expression env lhs_tp rhs;
        env, RT_Unit

      | STMT_MultiVarDecl {vars; init; loc} ->
        begin
          let ret_types=check_function_args env init in
          let expected = List.length ret_types in
          let actual = List.length vars in
          match init with Call {loc;callee;_} ->
          if expected!=actual then ErrorReporter.report_expected_function_returning_many_values ~loc ~id:callee ~expected ~actual else
          begin
            let rec aux vars types bindings= match vars,types with
              | []          , []       -> bindings
              | None::vars  , t::types -> aux vars types bindings
              | (Some (VarDecl{tp;loc;id} as v) )::vars, t::types ->
                let expected =te2type_decl env tp in
                if expected==t then aux vars types (v::bindings)
                else ErrorReporter.report_binding_type_mismatch ~loc ~expected ~id ~actual:t
              | _ -> failwith "Should not happen"
            in (List.fold_left add_var_decl env (aux vars ret_types [])),RT_Unit
          end
        end

      | STMT_Block body ->
        check_statement_block env body

      | STMT_Call call ->
        check_procedure_call env call;
        env, RT_Unit

      | STMT_If {cond;then_branch;else_branch; _} ->
        check_expression env TP_Bool cond;
        let if_ret = snd(check_statement env then_branch) in
        begin
          match else_branch with 
          | None -> env,RT_Unit
          | Some else_branch -> let else_ret =snd(check_statement env else_branch) in
            env, (if if_ret==RT_Void && else_ret==RT_Void then RT_Void else RT_Unit)
        end

      | STMT_Return {values;loc} ->
        begin
          match TypingEnvironment.get_return env with
          | None -> begin
                      match values with
                      | [] -> env,RT_Void
                      | _ -> ErrorReporter.report_procedure_cannot_return_value ~loc
                    end
          | Some ret ->(
            let expected = List.length ret
            and actual   = List.length values in
              if actual==0 then ErrorReporter.report_function_must_return_something ~loc
              else
              let rec aux rets values = match rets,values with
                | ([],[]) -> env, RT_Void
                | ((r::rs) , (v::vs)) -> let _=check_expression env r v in aux rs vs
                | ([],(x::xs)) -> ErrorReporter.report_bad_number_of_return_values ~loc ~expected ~actual
                | ((x::xs),[]) -> ErrorReporter.report_bad_number_of_return_values ~loc ~expected ~actual
              in aux ret values)

        end

      | STMT_VarDecl {var; init} ->
        begin
        match var with 
        | VarDecl {tp;_} ->
          let t=te2type_decl env tp in
            begin
              match init with
              | Some init -> check_expression env t init
              | None -> ()
            end;
          (add_var_decl env var),RT_Unit
        end

      | STMT_While {cond; body; _} ->
        check_expression env TP_Bool cond;
        let _ = check_statement env body in
        env,RT_Unit

    and check_statement_block env (STMTBlock {body; _}) =
        let rec aux env = function
          | []  -> env,RT_Unit
          | [x] -> check_statement env x
          | x::xs ->
          begin match check_statement env x with
                | (env,RT_Unit) -> aux env xs
                | (env,RT_Void) -> ErrorReporter.report_other_error ~loc:(location_of_statement x) ~descr:"Return not at the end of block - unreachable code"
          end
        in env,(snd (aux env body))

    (* --------------------------------------------------- *)
    (* Top-level funkcje *)

    let check_global_declaration env = function
      | GDECL_Function {formal_parameters; return_types; body; loc; id; _} ->
        (* Sprawdź wszystko *)
        let env = List.fold_left add_var_decl env formal_parameters in
        match body with
        | Some body -> begin
                      match check_statement_block  (TypingEnvironment.set_return env (List.map te2type_ret return_types)) body with
                        | _, RT_Unit -> if return_types!=[] then ErrorReporter.report_not_all_control_paths_return_value ~loc ~id else ()
                        | _, RT_Void -> ()
                        end
        | None -> ()

    let scan_global_declaration env = function
      | GDECL_Function {id; formal_parameters; return_types; loc; _} ->
         match TypingEnvironment.add id (ENVTP_Fn ((List.map (fun (VarDecl x) -> te2type_ret(x.tp)) formal_parameters),(List.map te2type_ret return_types)) ) env with
        | (env,true)  -> env
        | (env,false) -> ErrorReporter.report_shadows_previous_definition ~loc ~id

    let scan_module env (ModuleDefinition {global_declarations; _}) =
      List.fold_left scan_global_declaration env global_declarations

    let check_module env (ModuleDefinition {global_declarations; _}) =
      List.iter (check_global_declaration env) global_declarations

    (* --------------------------------------------------- *)
    (* Przetwórz moduł *)
    let process_module env mdef =
      (* Wpierw przeskanuj globalne definicje aby uzupełnić środowisko *)
      let env = scan_module env mdef in
      (* Zweryfikuj wszystko *)
      check_module env mdef

    let computation mdef = 
      (* Zaczynamy z pustym środowiskiem *)
      let env = TypingEnvironment.empty in
      process_module env mdef;
      node2type_map
  end

  (* --------------------------------------------------- *)
  (* Procedura wejściowa *)
  let check_module mdef =
    (* Stwórz instancję typecheckera i ją odpal *)
    let module M = Check() in
    M.ErrorReporter.wrap M.computation mdef 

end