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
    | _,   .univ i     => pure (.univ i)
    | 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
    | t,      u => pure (.app t u)

  partial def vFst : Val → EvalM Val
    | .pair a _ => pure a
    | t         => pure (.fst t)

  partial def vSnd : Val → EvalM Val
    | .pair _ b => pure b
    | t         => pure (.snd t)

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

partial def quote : Lvl → Val → 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 ← quote l t
      let qu ← quote l u
      pure (.app qt qu)
  | l, .fst t    => do
      let qt ← quote l t
      pure (.fst qt)
  | l, .snd t    => do
      let qt ← quote l t
      pure (.snd qt)
  | l, .lam c    => do
      let body ← cApp c (.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 (.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 (.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)
  | _, .univ i   => pure (.univ i)

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

partial def conv : Lvl → Val → Val → EvalM Bool
  | _, .univ i,     .univ j     => pure (i == j)
  | l, .pi a c,     .pi a' c'   =>
      andThen (conv l a a') fun _ => do
        let b  ← cApp c (.var l)
        let b' ← cApp c' (.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 (.var l)
        let b' ← cApp c' (.var l)
        conv (l + 1) b b'
  | l, .lam c,      .lam c'     =>
      do
        let body  ← cApp c (.var l)
        let body' ← cApp c' (.var l)
        conv (l + 1) body body'
  | l, .lam c,      t           =>
      do
        let body  ← cApp c (.var l)
        let apped ← vApp t (.var l)
        conv (l + 1) body apped
  | l, t,           .lam c      =>
      do
        let apped ← vApp t (.var l)
        let body  ← cApp c (.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
  | _, .var x,      .var y      => pure (x == y)
  | l, .app t u,    .app t' u'  =>
      andThen (conv l t t') fun _ => conv l u u'
  | l, .fst t,      .fst t'     => conv l t t'
  | l, .snd t,      .snd t'     => conv l t t'
  | _, _,           _           => pure false

end BidirTT