import BidirTT.Context
import BidirTT.Eval
import BidirTT.Pretty

namespace BidirTT

abbrev TCM := Except String

private def showTy (l : Lvl) (v : Val) : String :=
  match quote l v with
  | Except.ok tm     => prettyTm tm
  | Except.error err => s!"<ill-scoped value: {err}>"

mutual
  partial def check (cxt : Cxt) : Raw → Val → TCM Tm
    | .lam x t, .pi a c => do
        let bodyTy ← cApp c (.neu (.var cxt.lvl))
        let t' ← check (cxt.bind x a) t bodyTy
        pure (.lam t')
    | .pair t u, .sig a c => do
        let t' ← check cxt t a
        let vt ← eval cxt.env t'
        let bodyTy ← cApp c vt
        let u' ← check cxt u bodyTy
        pure (.pair t' u')
    | .refl, .id a t u => do
        if (← conv cxt.lvl t u) then
          pure .refl
        else
          throw s!"refl cannot inhabit {showTy cxt.lvl (.id a t u)} because the endpoints are not definitionally equal"
    | .letE x a t u, ty => do
        let (a', _) ← inferUniverse cxt a
        let va := eval cxt.env a'
        let va ← va
        let t' ← check cxt t va
        let vt ← eval cxt.env t'
        let u' ← check (cxt.define x vt va) u ty
        pure (.letE a' t' u')
    | r, ty => do
        let (t', ty') ← infer cxt r
        if (← sub cxt.lvl ty' ty) then
          pure t'
        else
          throw s!"type mismatch: expected {showTy cxt.lvl ty}, got {showTy cxt.lvl ty'}"

  partial def infer (cxt : Cxt) : Raw → TCM (Tm × Val)
    | .var x =>
        match cxt.lookup x with
        | some (i, a) => pure (.var i.val, a)
        | none        => throw s!"unknown variable {x}"
    | .nat => pure (.nat, .univ 0)
    | .zero => pure (.zero, .nat)
    | .succ t => do
        let t' ← check cxt t .nat
        pure (.succ t', .nat)
    | .natElim n motive z k ih s scrut => do
        let scrut' ← check cxt scrut .nat
        let vscrut ← eval cxt.env scrut'
        let (motiveBody', level) ← inferUniverse (cxt.bind n .nat) motive
        let motive' : Tm := .lam motiveBody'
        let vmotive ← eval cxt.env motive'
        let zTy ← vApp vmotive .zero
        let z' ← check cxt z zTy
        let kCxt := cxt.bind k .nat
        let ihTy ← vApp vmotive (.neu (.var cxt.lvl))
        let stepTy ← vApp vmotive (.succ (.neu (.var cxt.lvl)))
        let stepBody' ← check (kCxt.bind ih ihTy) s stepTy
        let step' : Tm := .lam (.lam stepBody')
        let resultTy ← vApp vmotive vscrut
        let _ := level
        pure (.natElim motive' z' step' scrut', resultTy)
    | .unit => pure (.unit, .univ 0)
    | .triv => pure (.triv, .unit)
    | .unitElim u motive t scrut => do
        let scrut' ← check cxt scrut .unit
        let vscrut ← eval cxt.env scrut'
        let (motiveBody', level) ← inferUniverse (cxt.bind u .unit) motive
        let motive' : Tm := .lam motiveBody'
        let vmotive ← eval cxt.env motive'
        let tTy ← vApp vmotive .triv
        let t' ← check cxt t tTy
        let resultTy ← vApp vmotive vscrut
        let _ := level
        pure (.unitElim motive' t' scrut', resultTy)
    | .empty => pure (.empty, .univ 0)
    | .emptyElim e motive scrut => do
        let scrut' ← check cxt scrut .empty
        let vscrut ← eval cxt.env scrut'
        let (motiveBody', level) ← inferUniverse (cxt.bind e .empty) motive
        let motive' : Tm := .lam motiveBody'
        let vmotive ← eval cxt.env motive'
        let resultTy ← vApp vmotive vscrut
        let _ := level
        pure (.emptyElim motive' scrut', resultTy)
    | .id a t u => do
        let (a', level) ← inferUniverse cxt a
        let va ← eval cxt.env a'
        let t' ← check cxt t va
        let u' ← check cxt u va
        pure (.id a' t' u', .univ level)
    | .refl => throw "cannot infer type of refl, use an annotation"
    | .idElim y p motive r target eq => do
        let (eq', eqTy) ← infer cxt eq
        match eqTy with
        | .id a x rhs => do
            let (target', targetTy) ← infer cxt target
            if !(← conv cxt.lvl targetTy a) then
              throw s!"idElim target has type {showTy cxt.lvl targetTy}, but the equality lives over {showTy cxt.lvl a}"
            let vtarget ← eval cxt.env target'
            if !(← conv cxt.lvl vtarget rhs) then
              throw s!"idElim target {showTy cxt.lvl vtarget} does not match the equality endpoint {showTy cxt.lvl rhs}"
            let eqVarTy : Val := .id a x (.neu (.var cxt.lvl))
            let (motiveBody', level) ← inferUniverse ((cxt.bind y a).bind p eqVarTy) motive
            let motive' : Tm := .lam (.lam motiveBody')
            let vmotive ← eval cxt.env motive'
            let reflTy ← vApp vmotive x
            let reflTy ← vApp reflTy .refl
            let r' ← check cxt r reflTy
            let veq ← eval cxt.env eq'
            let resultTy ← vApp vmotive vtarget
            let resultTy ← vApp resultTy veq
            let _ := level
            pure (.idElim motive' r' target' eq', resultTy)
        | _ => throw s!"expected Id type in idElim, got {showTy cxt.lvl eqTy}"
    | .univ i => pure (.univ i.normalise, .univ i.succ')
    | .app t u => do
        let (t', tty) ← infer cxt t
        match tty with
        | .pi a c => do
            let u' ← check cxt u a
            let vu ← eval cxt.env u'
            let bodyTy ← cApp c vu
            pure (.app t' u', bodyTy)
        | _ => throw s!"expected Pi type in application, got {showTy cxt.lvl tty}"
    | .fst t => do
        let (t', tty) ← infer cxt t
        match tty with
        | .sig a _ => pure (.fst t', a)
        | _ => throw s!"expected Sigma type in .1, got {showTy cxt.lvl tty}"
    | .snd t => do
        let (t', tty) ← infer cxt t
        match tty with
        | .sig _ c => do
            let vt ← eval cxt.env t'
            let fstv ← vFst vt
            let bodyTy ← cApp c fstv
            pure (.snd t', bodyTy)
        | _ => throw s!"expected Sigma type in .2, got {showTy cxt.lvl tty}"
    | .pi x a b => do
        let (a', i) ← inferUniverse cxt a
        let va ← eval cxt.env a'
        let (b', j) ← inferUniverse (cxt.bind x va) b
        pure (.pi a' b', .univ (i.max' j))
    | .sig x a b => do
        let (a', i) ← inferUniverse cxt a
        let va ← eval cxt.env a'
        let (b', j) ← inferUniverse (cxt.bind x va) b
        pure (.sig a' b', .univ (i.max' j))
    | .ann t a => do
        let (a', _) ← inferUniverse cxt a
        let va ← eval cxt.env a'
        let t' ← check cxt t va
        pure (t', va)
    | .letE x a t u => do
        let (a', _) ← inferUniverse cxt a
        let va ← eval cxt.env a'
        let t' ← check cxt t va
        let vt ← eval cxt.env t'
        let (u', uty) ← infer (cxt.define x vt va) u
        pure (.letE a' t' u', uty)
    | .lam _ _  => throw "cannot infer type of lambda, use an annotation"
    | .pair _ _ => throw "cannot infer type of pair, use an annotation"

  partial def inferUniverse (cxt : Cxt) (r : Raw) : TCM (Tm × Level) := do
    let (t, ty) ← infer cxt r
    match ty with
    | .univ level => pure (t, level)
    | _ => throw s!"expected a universe, got {showTy cxt.lvl ty}"
end

def checkTop (r : Raw) : TCM (Tm × Val) :=
  infer Cxt.empty r

end BidirTT