open Bwd
open Bwd.Infix

module A = Ast
module S = NbE.Syntax
module D = NbE.Domain

type env = {
  (* local context *)
  (* invariant: `tps`, `tms` and `names` all have length `size` *)
  tps : D.env;
  tms : D.env;
  names : Name.local bwd;
  size : int;

  (* top-level context *)
  toplvl : TopLevel.t;
}

module Eff = Algaeff.Reader.Make (struct type nonrec t = env end)

(* general helpers *)

let lookup (name : Name.t) : D.t * S.t =
  let env = Eff.read() in
  (* search through local context *)
  match Option.bind 
          (Name.to_local name) 
          (fun name -> Bwd.find_index ((=) name) env.names) with
  | Some ix ->
    let tp = Bwd.nth env.tps ix in
    (tp, S.Var ix)
  | None ->
    (* look up in top-level context *)
    match Yuujinchou.Trie.find_singleton name env.toplvl with
    | Some (Def { tp; tm }, _) ->
      (tp, S.Def (name, tm))
    | None -> Error.unbound_variable name

let bind ~(name : Name.local) ~(tp : D.t) f =
  let arg = D.var (Eff.read()).size in
  let update env = { 
    env with
    tps = env.tps <: tp;
    tms = env.tms <: arg;
    names = env.names <: name;
    size = env.size + 1;
  } in
  Eff.scope update (fun () -> f arg)

(* NbE helpers *)

let eval tm = NbE.eval ~env:(Eff.read()).tms tm

(* evaluate under the current environment augmented by `arg` *)
(* TODO: this is kind of inelegant, can we do better? *)
let eval_at arg = NbE.eval ~env:((Eff.read()).tms <: arg)

(* pretty-printing helpers *)

let pp_tm () =
  let names = (Eff.read()).names in
  fun fmt tm -> Pretty.pp ~names fmt tm

let pp_val () =
  let size = (Eff.read()).size in
  let pp_tm = pp_tm () in
  fun fmt v -> pp_tm fmt (NbE.quote ~size v)

(* main algorithm *)

type connective = [ `Pi | `Sigma ]

let rec check ~(tm : A.expr) ~(tp : D.t) : S.t =
  Error.tracef ?loc:tm.loc "when checking against the type @[%a@]" (pp_val()) tp @@ fun () ->
  match tm.value with
  | A.Pi ((name, base), fam) -> check_connective `Pi ~name:name.value ~base ~fam ~tp
  | A.Fun (base, fam) -> check_connective `Pi ~name:None ~base ~fam ~tp
  | A.Lam (name, body) -> begin match tp with
      | D.Pi (_, base, fam) -> 
        let body = bind ~name:name.value ~tp:base @@ fun arg ->
          let fib = NbE.inst_clo fam arg in
          check ~tm:body ~tp:fib in
        S.Lam (name.value, body)
      | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "a Pi type"
    end
  | A.Sg ((name, base), fam) -> check_connective `Sigma ~name:name.value ~base ~fam ~tp
  | A.Prod (base, fam) -> check_connective `Sigma ~name:None ~base ~fam ~tp
  | A.Pair (fst, snd) -> begin match tp with
      | D.Sg (_, base, fam) -> 
        let fst = check ~tm:fst ~tp:base in
        let fib = NbE.inst_clo fam (eval fst) in
        let snd = check ~tm:snd ~tp:fib in
        S.Pair (fst, snd)
      | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "a Sigma type"
    end
  | A.Type -> begin match tp with (* TODO type-in-type *)
      | D.Type -> S.Type
      | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "type"
    end
  | A.Bool -> begin match tp with
      | D.Type -> S.Bool
      | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "type"
    end
  | A.True -> begin match tp with
      | D.Bool -> S.True
      | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "bool"
    end
  | A.False -> begin match tp with
      | D.Bool -> S.False
      | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "bool"
    end
  | _ -> let (inferred_tp, tm) = infer tm in begin
      try NbE.equate ~size:((Eff.read()).size) inferred_tp tp with
      | NbE.Unequal ->
        Error.type_mismatch (pp_val()) tp (pp_val()) inferred_tp
    end;
    tm

and check_connective connective ~(name : Name.local) ~(base : A.expr) ~(fam : A.expr) ~(tp : D.t) =
  match tp with
  | D.Type ->
    let base = check ~tm:base ~tp in
    let fam = bind ~name:name ~tp:(eval base) @@ fun _ -> check ~tm:fam ~tp in
    begin match connective with
      | `Pi -> S.Pi (name, base, fam)
      | `Sigma -> S.Sg (name, base, fam)
    end
  | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "type"

and check_tp (tp : A.expr) = eval @@ check ~tp:D.Type ~tm:tp

and infer (tm : A.expr) : D.t * S.t =
  Error.tracef ?loc:tm.loc "when inferring the type" @@ fun () ->
  match tm.value with
  | A.Var name -> lookup name
  | A.Check (tm, tp) ->
    let tp = check_tp tp in
    let tm = check ~tp ~tm in
    (tp, tm)
  | A.App (fn, arg) -> begin match infer fn with
      | (D.Pi (_, base, fam), fn) ->
        let arg = check ~tm:arg ~tp:base in
        let tp = NbE.inst_clo fam (eval arg) in
        let tm = S.App (fn, arg) in
        (tp, tm)
      | (tp, _) -> Error.type_mismatch ?loc:fn.loc Fmt.string "a Pi type" (pp_val()) tp
    end
  | A.Fst p -> begin match infer p with
      | (D.Sg (_, base, _), p) -> (base, S.Fst p)
      | (tp, _) -> Error.type_mismatch ?loc:p.loc Fmt.string "a Sigma type" (pp_val()) tp
    end
  | A.Snd p -> begin match infer p with
      | (D.Sg (_, _, fam), p) ->
        let tp = NbE.inst_clo fam (eval (S.Fst p)) in
        let tm = S.Snd p in
        (tp, tm)
      | (tp, _) -> Error.type_mismatch ?loc:p.loc Fmt.string "a Sigma type" (pp_val()) tp
    end
  | A.BoolElim { motive_var; motive_body; true_case; false_case; scrut } ->
    let scrut = check ~tm:scrut ~tp:D.Bool in
    let motive = bind ~name:motive_var.value ~tp:D.Bool @@ fun _ ->
      check ~tm:motive_body ~tp:D.Type in
    let motive_true = eval_at D.True motive in
    let motive_false = eval_at D.False motive in
    let motive_scrut = eval_at (eval scrut) motive in
    let true_case = check ~tm:true_case ~tp:motive_true in
    let false_case = check ~tm:false_case ~tp:motive_false in
    let tm = S.BoolElim { motive_var = motive_var.value; motive; true_case; false_case; scrut } in
    (motive_scrut, tm)
  | _ -> Error.not_inferable ()

(* interface *)

let initial_env toplvl : env = {
  tps = Emp;
  tms = Emp;
  names = Emp;
  size = 0;
  toplvl;
}

let check_toplevel ~(toplvl : TopLevel.t) ~(tm : A.expr) ~(tp : D.t) =
  Eff.run ~env:(initial_env toplvl) @@ fun () -> check ~tm ~tp

let check_tp_toplevel ~(toplvl : TopLevel.t) (tp : A.expr) =
  Eff.run ~env:(initial_env toplvl) @@ fun () -> check_tp tp

let infer_toplevel ~(toplvl : TopLevel.t) (tm : A.expr) =
  Eff.run ~env:(initial_env toplvl) @@ fun () -> infer tm