(* 타입 검사.

   범위: 이 모듈 안에서 아는 것만 검사한다. 외부 이름은 TUnknown이 되어
   무엇과도 맞는다 — 모르는 것을 틀렸다고 말하지 않는다.

   제네릭 해소는 호출 지점의 지역 unification이다. 함수 하나를 넘어가지
   않으므로 전역 추론이 아니다. *)

open Ast
module T = Types

type error = { pos : Token.pos; msg : string }

type scheme = {
  s_gen : string list; (* 타입 파라미터 이름 (effect 파라미터는 제외) *)
  s_params : T.t list;
  s_ret : T.t;
}

type env = {
  structs : (string, string list * (string * T.t) list) Hashtbl.t;
  enums : (string, string list * (string * T.t list) list) Hashtbl.t;
  caps : (string, (string * scheme) list) Hashtbl.t;
  fns : (string, scheme) Hashtbl.t;
  consts : (string, T.t) Hashtbl.t;
  ctors : (string, string) Hashtbl.t; (* variant -> enum *)
  mutable locals : (string * T.t) list list;
  mutable ret : T.t; (* 현재 함수의 선언된 반환 타입 *)
  mutable errors : error list;
}

let err env pos msg = env.errors <- { pos; msg } :: env.errors

let mismatch env pos expected got what =
  err env pos
    (Printf.sprintf "%s: %s이(가) 필요한데 %s입니다" what (T.show expected) (T.show got))

let push env = env.locals <- [] :: env.locals
let pop env = match env.locals with _ :: r -> env.locals <- r | [] -> ()

let bind env n t =
  match env.locals with
  | s :: r -> env.locals <- ((n, t) :: s) :: r
  | [] -> env.locals <- [ [ (n, t) ] ]

let lookup env n =
  let rec go = function
    | [] -> None
    | s :: r -> (
        match List.assoc_opt n s with Some t -> Some t | None -> go r)
  in
  go env.locals

(* ------------------------------------------------------------------ *)
(* Ast.ty -> Types.t                                                   *)
(* ------------------------------------------------------------------ *)

let rec conv env (gen : string list) (t : Ast.ty) : T.t =
  match t with
  | T_named { name; args; _ } -> (
      let args =
        List.filter_map
          (function TA_ty t -> Some (conv env gen t) | TA_eff _ -> None)
          args
      in
      if List.mem name gen then T.TVar name
      else
        match name with
        | "Int" -> T.TInt
        | "Bool" -> T.TBool
        | "String" -> T.TString
        | "Unit" -> T.TUnit
        | "List" | "Option" | "Result" -> T.TCon (name, args)
        | _ ->
            if
              Hashtbl.mem env.structs name
              || Hashtbl.mem env.enums name || Hashtbl.mem env.caps name
            then T.TCon (name, args)
            else T.TUnknown)
  | T_fn { affine; params; ret; _ } ->
      T.TFn
        {
          affine;
          params = List.map (conv env gen) params;
          ret = (match ret with None -> T.TUnit | Some r -> conv env gen r);
        }

let scheme_of env (d : fn_decl) : scheme =
  let gen =
    List.filter_map
      (fun g -> if g.gp_effect then None else Some g.gp_name)
      d.fn_gen
  in
  {
    s_gen = gen;
    s_params = List.map (fun p -> conv env gen p.p_ty) d.fn_params;
    s_ret = (match d.fn_ret with None -> T.TUnit | Some r -> conv env gen r);
  }

(* 호출 지점 인스턴스화: 타입 파라미터마다 새 미지수 *)
let instantiate (s : scheme) =
  let sub = List.map (fun v -> (v, T.fresh ())) s.s_gen in
  (List.map (T.subst sub) s.s_params, T.subst sub s.s_ret)

let instantiate_with (s : scheme) (args : T.t list) =
  let sub = List.map2 (fun v a -> (v, a)) s.s_gen args in
  (List.map (T.subst sub) s.s_params, T.subst sub s.s_ret)

(* ------------------------------------------------------------------ *)
(* 내장 생성자                                                          *)
(* ------------------------------------------------------------------ *)

let builtin_ctor = function
  | "Ok" ->
      let a = T.fresh () and b = T.fresh () in
      Some ([ a ], T.TCon ("Result", [ a; b ]))
  | "Err" ->
      let a = T.fresh () and b = T.fresh () in
      Some ([ b ], T.TCon ("Result", [ a; b ]))
  | "Some" ->
      let a = T.fresh () in
      Some ([ a ], T.TCon ("Option", [ a ]))
  | _ -> None

(* ------------------------------------------------------------------ *)
(* 식                                                                  *)
(* ------------------------------------------------------------------ *)

let rec infer env (e : expr) : T.t =
  match e with
  | E_lit (L_int _, _) -> T.TInt
  | E_lit (L_str _, _) -> T.TString
  | E_lit (L_bool _, _) -> T.TBool
  | E_ident (n, _) -> (
      match lookup env n with
      | Some t -> t
      | None -> (
          match Hashtbl.find_opt env.consts n with
          | Some t -> t
          | None -> (
              match n with
              | "unit" -> T.TUnit
              | "None" -> T.TCon ("Option", [ T.fresh () ])
              | _ -> (
                  match Hashtbl.find_opt env.fns n with
                  | Some s ->
                      let params, ret = instantiate s in
                      T.TFn { affine = false; params; ret }
                  | None -> (
                      match Hashtbl.find_opt env.ctors n with
                      | Some enum -> nullary_ctor env enum n
                      | None -> T.TUnknown)))))
  | E_list (xs, pos) ->
      let elem = T.fresh () in
      List.iter
        (fun x ->
          let t = infer env x in
          if not (T.unify elem t) then
            mismatch env pos elem t "리스트 원소의 타입이 서로 다릅니다")
        xs;
      T.TCon ("List", [ elem ])
  | E_struct { name; fields; pos } -> infer_struct env name fields pos
  | E_closure c -> infer_closure env c None
  | E_if { cond; then_; else_; pos } -> (
      let c = infer env cond in
      if not (T.unify c T.TBool) then mismatch env pos T.TBool c "if의 조건";
      let t1 = infer_block env then_ in
      match else_ with
      | None ->
          if not (T.unify t1 T.TUnit) then
            err env pos "else가 없는 if의 본문은 값을 남길 수 없습니다";
          T.TUnit
      | Some e2 ->
          let t2 = infer env e2 in
          if not (T.unify t1 t2) then mismatch env pos t1 t2 "if의 두 분기 타입이 다릅니다";
          t1)
  | E_match { scrutinee; arms; pos } ->
      let s = infer env scrutinee in
      let result = T.fresh () in
      List.iter
        (fun a ->
          push env;
          check_pattern env s a.arm_pat;
          let t = infer env a.arm_body in
          if not (T.unify result t) then
            mismatch env a.arm_pos result t "match 팔의 타입이 서로 다릅니다";
          pop env)
        arms;
      if arms = [] then err env pos "match에 팔이 없습니다";
      result
  | E_scope { name; parent; body; pos } ->
      (match lookup env parent with
      | Some t when not (T.unify t (T.TCon ("TaskScope", []))) ->
          if T.resolve t <> T.TUnknown then
            mismatch env pos (T.TCon ("TaskScope", [])) t "scope의 부모"
      | _ -> ());
      push env;
      bind env name (T.TCon ("TaskScope", []));
      let t = infer_block env body in
      pop env;
      t
  | E_block b ->
      push env;
      let t = infer_block env b in
      pop env;
      t
  | E_call { callee; args; pos } -> infer_call env callee args pos
  | E_field { obj; name; pos } -> infer_field env obj name pos
  | E_inst { callee; args; pos } -> infer_inst env callee args pos
  | E_try { inner; pos } -> (
      let t = infer env inner in
      match T.resolve t with
      | T.TUnknown -> T.TUnknown
      | T.TCon ("Result", [ ok; _ ]) ->
          (match T.resolve env.ret with
          | T.TCon ("Result", _) | T.TUnknown -> ()
          | r ->
              err env pos
                (Printf.sprintf "?는 Result를 반환하는 함수 안에서만 쓸 수 있습니다 (현재 반환 타입 %s)"
                   (T.show r)));
          ok
      | other ->
          err env pos
            (Printf.sprintf "?는 Result에만 쓸 수 있습니다 (%s에 쓰였습니다)" (T.show other));
          T.TUnknown)
  | E_unary { op; operand; pos } ->
      let t = infer env operand in
      let want = match op with U_not -> T.TBool | U_neg -> T.TInt in
      if not (T.unify t want) then mismatch env pos want t "단항 연산자의 피연산자";
      want
  | E_binary { op; lhs; rhs; pos } -> (
      let a = infer env lhs and b = infer env rhs in
      match op with
      | B_or | B_and ->
          if not (T.unify a T.TBool) then mismatch env pos T.TBool a "논리 연산자";
          if not (T.unify b T.TBool) then mismatch env pos T.TBool b "논리 연산자";
          T.TBool
      | B_add | B_sub | B_mul | B_div | B_rem ->
          if not (T.unify a T.TInt) then mismatch env pos T.TInt a "산술 연산자";
          if not (T.unify b T.TInt) then mismatch env pos T.TInt b "산술 연산자";
          T.TInt
      | B_lt | B_le | B_gt | B_ge ->
          if not (T.unify a T.TInt) then mismatch env pos T.TInt a "비교 연산자";
          if not (T.unify b T.TInt) then mismatch env pos T.TInt b "비교 연산자";
          T.TBool
      | B_eq | B_ne ->
          if not (T.unify a b) then mismatch env pos a b "같은 타입끼리만 비교할 수 있습니다";
          T.TBool)

and nullary_ctor env enum name =
  match Hashtbl.find_opt env.enums enum with
  | None -> T.TUnknown
  | Some (gen, variants) -> (
      let sub = List.map (fun v -> (v, T.fresh ())) gen in
      match List.assoc_opt name variants with
      | Some [] -> T.TCon (enum, List.map snd sub)
      | _ -> T.TCon (enum, List.map snd sub))

and infer_struct env name fields pos =
  match Hashtbl.find_opt env.structs name with
  | None ->
      List.iter (fun (_, e) -> ignore (infer env e)) fields;
      T.TUnknown
  | Some (gen, decl_fields) ->
      let sub = List.map (fun v -> (v, T.fresh ())) gen in
      List.iter
        (fun (fname, fe) ->
          match List.assoc_opt fname decl_fields with
          | None -> err env pos (Printf.sprintf "%s에 %s 필드가 없습니다" name fname)
          | Some ft ->
              let want = T.subst sub ft in
              let got = infer env fe in
              if not (T.unify want got) then
                mismatch env pos want got (Printf.sprintf "%s.%s 필드" name fname))
        fields;
      List.iter
        (fun (fname, _) ->
          if not (List.mem_assoc fname fields) then
            err env pos (Printf.sprintf "%s의 %s 필드가 빠졌습니다" name fname))
        decl_fields;
      T.TCon (name, List.map snd sub)

and infer_closure env (c : closure) (expected : T.t option) =
  let expected_params, expected_ret =
    match Option.map T.resolve expected with
    | Some (T.TFn { params; ret; _ })
      when List.length params = List.length c.cl_params ->
        (List.map Option.some params, Some ret)
    | _ -> (List.map (fun _ -> None) c.cl_params, None)
  in
  push env;
  let param_tys =
    List.map2
      (fun (n, ann) exp ->
        let t =
          match ann with
          | Some a -> conv env [] a
          | None -> ( match exp with Some t -> t | None -> T.TUnknown)
        in
        bind env n t;
        t)
      c.cl_params expected_params
  in
  let declared_ret = Option.map (conv env []) c.cl_ret in
  let saved = env.ret in
  env.ret <-
    (match declared_ret with
    | Some t -> t
    | None -> ( match expected_ret with Some t -> t | None -> T.TUnknown));
  let body = infer_block env c.cl_body in
  (match declared_ret with
  | Some t when not (T.unify t body) -> mismatch env c.cl_pos t body "클로저의 반환"
  | _ -> ());
  let ret = match declared_ret with Some t -> t | None -> body in
  env.ret <- saved;
  pop env;
  T.TFn { affine = false; params = param_tys; ret }

and infer_call env callee args pos =
  let fn_ty =
    match callee with
    | E_ident (n, _) when lookup env n = None -> (
        match Hashtbl.find_opt env.ctors n with
        | Some enum -> Some (ctor_fn env enum n)
        | None -> (
            match builtin_ctor n with
            | Some (params, ret) -> Some (T.TFn { affine = false; params; ret })
            | None -> (
                match Hashtbl.find_opt env.fns n with
                | Some s ->
                    let params, ret = instantiate s in
                    Some (T.TFn { affine = false; params; ret })
                | None -> None)))
    | _ -> (
        match T.resolve (infer env callee) with
        | T.TFn _ as t -> Some t
        | _ -> None)
  in
  match fn_ty with
  | None ->
      List.iter (fun a -> ignore (infer env a)) args;
      T.TUnknown
  | Some (T.TFn { params; ret; _ }) ->
      if List.length params <> List.length args then (
        err env pos
          (Printf.sprintf "인자 %d개가 필요한데 %d개가 주어졌습니다" (List.length params)
             (List.length args));
        List.iter (fun a -> ignore (infer env a)) args)
      else
        List.iter2
          (fun p a ->
            let got =
              match a with
              | E_closure c -> infer_closure env c (Some p)
              | _ -> infer env a
            in
            if not (T.unify p got) then mismatch env pos p got "인자")
          params args;
      ret
  | Some _ -> T.TUnknown

and ctor_fn env enum name =
  match Hashtbl.find_opt env.enums enum with
  | None -> T.TUnknown
  | Some (gen, variants) ->
      let sub = List.map (fun v -> (v, T.fresh ())) gen in
      let params =
        match List.assoc_opt name variants with
        | Some ts -> List.map (T.subst sub) ts
        | None -> []
      in
      T.TFn { affine = false; params; ret = T.TCon (enum, List.map snd sub) }

and infer_field env obj name pos =
  let t = infer env obj in
  match T.resolve t with
  | T.TUnknown -> T.TUnknown
  | T.TCon (cname, args) -> (
      match Hashtbl.find_opt env.structs cname with
      | Some (gen, fields) -> (
          let sub = List.map2 (fun v a -> (v, a)) gen (adjust gen args) in
          match List.assoc_opt name fields with
          | Some ft -> T.subst sub ft
          | None ->
              err env pos (Printf.sprintf "%s에 %s 필드가 없습니다" cname name);
              T.TUnknown)
      | None -> (
          match Hashtbl.find_opt env.caps cname with
          | Some methods -> (
              match List.assoc_opt name methods with
              | Some s ->
                  let params, ret = instantiate s in
                  T.TFn { affine = false; params; ret }
              | None ->
                  err env pos
                    (Printf.sprintf "capability %s에 %s 메서드가 없습니다" cname name);
                  T.TUnknown)
          | None -> T.TUnknown))
  | other ->
      err env pos (Printf.sprintf "%s에는 필드가 없습니다" (T.show other));
      T.TUnknown

and adjust gen args =
  let n = List.length gen in
  let rec take k xs =
    if k = 0 then []
    else
      match xs with
      | [] -> T.fresh () :: take (k - 1) []
      | x :: r -> x :: take (k - 1) r
  in
  take n args

and infer_inst env callee args pos =
  let tys =
    List.filter_map
      (function TA_ty t -> Some (conv env [] t) | TA_eff _ -> None)
      args
  in
  match callee with
  | E_ident (n, _) when lookup env n = None -> (
      match Hashtbl.find_opt env.fns n with
      | None -> T.TUnknown
      | Some s ->
          if List.length tys <> List.length s.s_gen then (
            err env pos
              (Printf.sprintf "타입 인자 %d개가 필요한데 %d개가 주어졌습니다"
                 (List.length s.s_gen) (List.length tys));
            T.TUnknown)
          else
            let params, ret = instantiate_with s tys in
            T.TFn { affine = false; params; ret })
  | _ -> T.TUnknown

and check_pattern env (scrutinee : T.t) (p : pattern) =
  match p with
  | P_wild _ -> ()
  | P_lit (l, pos) ->
      let t =
        match l with
        | L_int _ -> T.TInt
        | L_str _ -> T.TString
        | L_bool _ -> T.TBool
      in
      if not (T.unify scrutinee t) then mismatch env pos scrutinee t "패턴의 리터럴"
  | P_bind (n, _) -> (
      match Hashtbl.find_opt env.ctors n with
      | Some enum ->
          check_ctor env scrutinee enum n [] Token.{ line = 0; col = 0 }
      | None -> bind env n scrutinee)
  | P_ctor { name; args; pos } -> (
      match Hashtbl.find_opt env.ctors name with
      | Some enum -> check_ctor env scrutinee enum name args pos
      | None -> List.iter (check_pattern env T.TUnknown) args)

and check_ctor env scrutinee enum name args pos =
  match Hashtbl.find_opt env.enums enum with
  | None -> ()
  | Some (gen, variants) ->
      let sub = List.map (fun v -> (v, T.fresh ())) gen in
      let ety = T.TCon (enum, List.map snd sub) in
      if not (T.unify scrutinee ety) then
        mismatch env pos scrutinee ety "패턴이 match 대상과 다른 타입입니다";
      let fields =
        match List.assoc_opt name variants with Some ts -> ts | None -> []
      in
      if List.length fields = List.length args then
        List.iter2
          (fun ft ap -> check_pattern env (T.subst sub ft) ap)
          fields args

(* ------------------------------------------------------------------ *)
(* 문과 블록                                                           *)
(* ------------------------------------------------------------------ *)

and infer_block env (b : block) : T.t =
  let rec go = function
    | [] -> T.TUnit
    | [ S_expr e ] -> infer env e
    | s :: rest ->
        check_stmt env s;
        go rest
  in
  go b.stmts

and check_stmt env = function
  | S_let { pat; ty; value; pos; _ } ->
      let declared = Option.map (conv env []) ty in
      let got =
        match (value, declared) with
        | E_closure c, Some t -> infer_closure env c (Some t)
        | _ -> infer env value
      in
      let t =
        match declared with
        | None -> got
        | Some d ->
            if not (T.unify d got) then mismatch env pos d got "let의 타입 주석";
            d
      in
      check_pattern env t pat
  | S_return { value; pos } ->
      let got = match value with None -> T.TUnit | Some e -> infer env e in
      if not (T.unify env.ret got) then mismatch env pos env.ret got "return의 값"
  | S_assign { place; value; pos } ->
      let p = infer env place in
      let v = infer env value in
      if not (T.unify p v) then mismatch env pos p v "대입"
  | S_expr e -> ignore (infer env e)

(* ------------------------------------------------------------------ *)
(* 모듈                                                                *)
(* ------------------------------------------------------------------ *)

let check_fn env (d : fn_decl) =
  match d.fn_body with
  | None -> ()
  | Some body ->
      let gen =
        List.filter_map
          (fun g -> if g.gp_effect then None else Some g.gp_name)
          d.fn_gen
      in
      push env;
      List.iter (fun p -> bind env p.p_name (conv env gen p.p_ty)) d.fn_params;
      let declared =
        match d.fn_ret with None -> T.TUnit | Some r -> conv env gen r
      in
      env.ret <- declared;
      let got = infer_block env body in
      if not (T.unify declared got) then
        mismatch env d.fn_pos declared got
          (Printf.sprintf "%s의 본문이 남기는 값" d.fn_name);
      pop env

let check (m : modul) : error list =
  let env =
    {
      structs = Hashtbl.create 16;
      enums = Hashtbl.create 16;
      caps = Hashtbl.create 16;
      fns = Hashtbl.create 16;
      consts = Hashtbl.create 16;
      ctors = Hashtbl.create 16;
      locals = [];
      ret = T.TUnit;
      errors = [];
    }
  in
  (* 1차: 타입과 생성자 이름부터 (선언 순서에 의존하지 않는다) *)
  List.iter
    (fun it ->
      match it with
      | I_struct { name; gen; _ } ->
          Hashtbl.replace env.structs name
            (List.map (fun g -> g.gp_name) gen, [])
      | I_enum { name; gen; variants; _ } ->
          Hashtbl.replace env.enums name (List.map (fun g -> g.gp_name) gen, []);
          List.iter (fun v -> Hashtbl.replace env.ctors v.v_name name) variants
      | I_capability { name; _ } -> Hashtbl.replace env.caps name []
      | _ -> ())
    m.items;
  (* 2차: 본문을 채운다 *)
  List.iter
    (fun it ->
      match it with
      | I_struct { name; gen; fields; _ } ->
          let g = List.map (fun x -> x.gp_name) gen in
          Hashtbl.replace env.structs name
            (g, List.map (fun f -> (f.f_name, conv env g f.f_ty)) fields)
      | I_enum { name; gen; variants; _ } ->
          let g = List.map (fun x -> x.gp_name) gen in
          Hashtbl.replace env.enums name
            ( g,
              List.map
                (fun v -> (v.v_name, List.map (conv env g) v.v_args))
                variants )
      | I_capability { name; methods; _ } ->
          Hashtbl.replace env.caps name
            (List.map (fun d -> (d.fn_name, scheme_of env d)) methods)
      | I_fn { decl; _ } ->
          Hashtbl.replace env.fns decl.fn_name (scheme_of env decl)
      | I_const { name; ty; _ } ->
          Hashtbl.replace env.consts name (conv env [] ty)
      | _ -> ())
    m.items;
  (* 3차: 본문 검사 *)
  List.iter
    (fun it ->
      env.locals <- [];
      match it with
      | I_fn { decl; _ } -> check_fn env decl
      | I_const { ty; value; pos; _ } ->
          let want = conv env [] ty in
          let got = infer env value in
          if not (T.unify want got) then mismatch env pos want got "상수의 값"
      | _ -> ())
    m.items;
  List.sort
    (fun a b ->
      compare
        (a.pos.Token.line, a.pos.Token.col)
        (b.pos.Token.line, b.pos.Token.col))
    (List.rev env.errors)