import BidirTT.Value

namespace BidirTT

abbrev EvalM := Except String

mutual
  partial def eval : Env → Tm → EvalM Val
    | env, .var i =>
        match env[i]? with
        | some v => pure v
        | none =>
            throw s!"bad de Bruijn index {i} in environment of size {env.length}"
    | env, .lam t      => pure (.lam (.mk env t))
    | env, .app t u    => do
        let vt ← eval env t
        let vu ← eval env u
        vApp vt vu
    | env, .pi a b     => do
        let va ← eval env a
        pure (.pi va (.mk env b))
    | env, .sig a b    => do
        let va ← eval env a
        pure (.sig va (.mk env b))
    | env, .pair t u   => do
        let vt ← eval env t
        let vu ← eval env u
        pure (.pair vt vu)
    | env, .fst t      => do
        let vt ← eval env t
        vFst vt
    | env, .snd t      => do
        let vt ← eval env t
        vSnd vt
    | _,   .nat        => pure .nat
    | _,   .zero       => pure .zero
    | env, .succ t     => do
        let vt ← eval env t
        pure (.succ vt)
    | env, .natElim m z s n => do
        let vm ← eval env m
        let vz ← eval env z
        let vs ← eval env s
        let vn ← eval env n
        vNatElim vm vz vs vn
    | _,   .unit       => pure .unit
    | _,   .triv       => pure .triv
    | env, .unitElim m t u => do
        let vm ← eval env m
        let vt ← eval env t
        let vu ← eval env u
        vUnitElim vm vt vu
    | _,   .empty      => pure .empty
    | env, .emptyElim m e => do
        let vm ← eval env m
        let ve ← eval env e
        vEmptyElim vm ve
    | env, .id a t u   => do
        let va ← eval env a
        let vt ← eval env t
        let vu ← eval env u
        pure (.id va vt vu)
    | _,   .refl       => pure .refl
    | env, .idElim m r y p => do
        let vm ← eval env m
        let vr ← eval env r
        let vy ← eval env y
        let vp ← eval env p
        vIdElim vm vr vy vp
    | _,   .univ i     => pure (.univ i.normalise)
    | env, .letE _ t u => do
        let vt ← eval env t
        eval (vt :: env) u

  partial def vApp : Val → Val → EvalM Val
    | .lam c,   u => cApp c u
    | .neu t,   u => pure (.neu (.app t u))
    | _,        _ => throw "bad application head during evaluation"

  partial def vFst : Val → EvalM Val
    | .pair a _ => pure a
    | .neu t     => pure (.neu (.fst t))
    | _          => throw "bad fst projection during evaluation"

  partial def vSnd : Val → EvalM Val
    | .pair _ b => pure b
    | .neu t    => pure (.neu (.snd t))
    | _         => throw "bad snd projection during evaluation"

  partial def vNatElim : Val → Val → Val → Val → EvalM Val
    | _, z, _, .zero   => pure z
    | m, z, s, .succ n => do
        let ih ← vNatElim m z s n
        let step ← vApp s n
        vApp step ih
    | m, z, s, .neu n  => pure (.neu (.natElim m z s n))
    | _, _, _, _       => throw "bad Nat eliminand during evaluation"

  partial def vUnitElim : Val → Val → Val → EvalM Val
    | _, t, .triv  => pure t
    | m, t, .neu u => pure (.neu (.unitElim m t u))
    | _, _, _      => throw "bad Unit eliminand during evaluation"

  partial def vEmptyElim : Val → Val → EvalM Val
    | m, .neu e => pure (.neu (.emptyElim m e))
    | _, _      => throw "bad Empty eliminand during evaluation"

  partial def vIdElim : Val → Val → Val → Val → EvalM Val
    | _, r, _, .refl => pure r
    | m, r, y, .neu p => pure (.neu (.idElim m r y p))
    | _, _, _, _      => throw "bad Id eliminand during evaluation"

  partial def cApp : Closure → Val → EvalM Val
    | .mk env body, v => eval (v :: env) body
end

mutual
  partial def quoteNeutral : Lvl → Neutral → EvalM Tm
    | l, .var x =>
        if x < l then
          pure (.var (l - x - 1))
        else
          throw s!"bad level {x} while quoting at level {l}"
    | l, .app t u => do
        let qt ← quoteNeutral l t
        let qu ← quote l u
        pure (.app qt qu)
    | l, .fst t => do
        let qt ← quoteNeutral l t
        pure (.fst qt)
    | l, .snd t => do
        let qt ← quoteNeutral l t
        pure (.snd qt)
    | l, .natElim m z s n => do
        let qm ← quote l m
        let qz ← quote l z
        let qs ← quote l s
        let qn ← quoteNeutral l n
        pure (.natElim qm qz qs qn)
    | l, .unitElim m t u => do
        let qm ← quote l m
        let qt ← quote l t
        let qu ← quoteNeutral l u
        pure (.unitElim qm qt qu)
    | l, .emptyElim m e => do
        let qm ← quote l m
        let qe ← quoteNeutral l e
        pure (.emptyElim qm qe)
    | l, .idElim m r y p => do
        let qm ← quote l m
        let qr ← quote l r
        let qy ← quote l y
        let qp ← quoteNeutral l p
        pure (.idElim qm qr qy qp)

  partial def quote : Lvl → Val → EvalM Tm
    | l, .neu n    => quoteNeutral l n
    | l, .lam c    => do
        let body ← cApp c (.neu (.var l))
        let qb ← quote (l + 1) body
        pure (.lam qb)
    | l, .pi a c   => do
        let qa ← quote l a
        let body ← cApp c (.neu (.var l))
        let qb ← quote (l + 1) body
        pure (.pi qa qb)
    | l, .sig a c  => do
        let qa ← quote l a
        let body ← cApp c (.neu (.var l))
        let qb ← quote (l + 1) body
        pure (.sig qa qb)
    | l, .pair a b => do
        let qa ← quote l a
        let qb ← quote l b
        pure (.pair qa qb)
    | _, .nat      => pure .nat
    | _, .zero     => pure .zero
    | l, .succ t   => do
        let qt ← quote l t
        pure (.succ qt)
    | _, .unit     => pure .unit
    | _, .triv     => pure .triv
    | _, .empty    => pure .empty
    | l, .id a t u => do
        let qa ← quote l a
        let qt ← quote l t
        let qu ← quote l u
        pure (.id qa qt qu)
    | _, .refl     => pure .refl
    | _, .univ i   => pure (.univ i.normalise)
end

private def andThen (lhs : EvalM Bool) (rhs : Unit → EvalM Bool) : EvalM Bool := do
  if (← lhs) then
    rhs ()
  else
    pure false

mutual
  partial def convNeutral : Lvl → Neutral → Neutral → EvalM Bool
    | _, .var x, .var y => pure (x == y)
    | l, .app t u, .app t' u' =>
        andThen (convNeutral l t t') fun _ => conv l u u'
    | l, .fst t, .fst t' =>
        convNeutral l t t'
    | l, .snd t, .snd t' =>
        convNeutral l t t'
    | l, .natElim m z s n, .natElim m' z' s' n' =>
        andThen (conv l m m') fun _ => do
          let sameZ ← conv l z z'
          if sameZ then
            let sameS ← conv l s s'
            if sameS then
              convNeutral l n n'
            else
              pure false
          else
            pure false
    | l, .unitElim m t u, .unitElim m' t' u' =>
        andThen (conv l m m') fun _ => do
          let sameT ← conv l t t'
          if sameT then
            convNeutral l u u'
          else
            pure false
    | l, .emptyElim m e, .emptyElim m' e' =>
        andThen (conv l m m') fun _ => convNeutral l e e'
    | l, .idElim m r y p, .idElim m' r' y' p' =>
        andThen (conv l m m') fun _ => do
          let sameR ← conv l r r'
          if sameR then
            let sameY ← conv l y y'
            if sameY then
              convNeutral l p p'
            else
              pure false
          else
            pure false
    | _, _, _ => pure false

  partial def conv : Lvl → Val → Val → EvalM Bool
    | _, .univ i,     .univ j     => Level.eqv i j
    | l, .pi a c,     .pi a' c'   =>
        andThen (conv l a a') fun _ => do
          let b  ← cApp c (.neu (.var l))
          let b' ← cApp c' (.neu (.var l))
          conv (l + 1) b b'
    | l, .sig a c,    .sig a' c'  =>
        andThen (conv l a a') fun _ => do
          let b  ← cApp c (.neu (.var l))
          let b' ← cApp c' (.neu (.var l))
          conv (l + 1) b b'
    | _, .nat,        .nat        => pure true
    | _, .zero,       .zero       => pure true
    | l, .succ n,     .succ n'    => conv l n n'
    | _, .unit,       .unit       => pure true
    | _, .triv,       .triv       => pure true
    | _, .empty,      .empty      => pure true
    | l, .id a t u,   .id a' t' u' =>
        andThen (conv l a a') fun _ => do
          let sameT ← conv l t t'
          if sameT then
            conv l u u'
          else
            pure false
    | _, .refl,       .refl       => pure true
    | l, .lam c,      .lam c'     => do
        let body  ← cApp c (.neu (.var l))
        let body' ← cApp c' (.neu (.var l))
        conv (l + 1) body body'
    | l, .lam c,      t           => do
        let body  ← cApp c (.neu (.var l))
        let apped ← vApp t (.neu (.var l))
        conv (l + 1) body apped
    | l, t,           .lam c      => do
        let apped ← vApp t (.neu (.var l))
        let body  ← cApp c (.neu (.var l))
        conv (l + 1) apped body
    | l, .pair a b,   .pair a' b' =>
        andThen (conv l a a') fun _ => conv l b b'
    | l, .pair a b,   p           =>
        andThen
          (do
            let fstp ← vFst p
            conv l a fstp)
          fun _ => do
            let sndp ← vSnd p
            conv l b sndp
    | l, p,           .pair a b   =>
        andThen
          (do
            let fstp ← vFst p
            conv l fstp a)
          fun _ => do
            let sndp ← vSnd p
            conv l sndp b
    | l, .neu n,      .neu n'     => convNeutral l n n'
    | _, _,           _           => pure false
end

partial def sub : Lvl → Val → Val → EvalM Bool
  | _, .univ i,     .univ j     => Level.leq i j
  | l, .pi a c,     .pi a' c'   =>
      andThen (sub l a' a) fun _ => do
        let b  ← cApp c (.neu (.var l))
        let b' ← cApp c' (.neu (.var l))
        sub (l + 1) b b'
  | l, .sig a c,    .sig a' c'  =>
      andThen (sub l a a') fun _ => do
        let b  ← cApp c (.neu (.var l))
        let b' ← cApp c' (.neu (.var l))
        sub (l + 1) b b'
  | l, t,           t'          => conv l t t'

end BidirTT