import BidirTT.Value

namespace BidirTT

abbrev CxtEntry := Name × Val

structure Cxt where
  env   : Env
  types : List CxtEntry
  lvl   : Lvl
  envWf : env.length = lvl
  typesWf : types.length = lvl

def Cxt.empty : Cxt := ⟨[], [], 0, rfl, rfl⟩

instance : Inhabited Cxt := ⟨Cxt.empty⟩

def Cxt.bind (cxt : Cxt) (x : Name) (a : Val) : Cxt :=
  { env   := .neu (.var cxt.lvl) :: cxt.env
  , types := (x, a) :: cxt.types
  , lvl   := cxt.lvl + 1
  , envWf := by simp [cxt.envWf]
  , typesWf := by simp [cxt.typesWf] }

def Cxt.define (cxt : Cxt) (x : Name) (v a : Val) : Cxt :=
  { env   := v :: cxt.env
  , types := (x, a) :: cxt.types
  , lvl   := cxt.lvl + 1
  , envWf := by simp [cxt.envWf]
  , typesWf := by simp [cxt.typesWf] }

private def castFin {n m : Nat} (h : n = m) (i : Fin n) : Fin m :=
  ⟨i.val, h ▸ i.isLt⟩

private def lookupGo (x : Name) : (entries : List CxtEntry) → Option (Fin entries.length × Val)
  | [] => none
  | (y, a) :: rest =>
      if x == y then
        some (⟨0, by simp⟩, a)
      else
        match lookupGo x rest with
        | some (i, a) =>
            some (⟨i.val + 1, by simp [i.isLt]⟩, a)
        | none =>
            none

def Cxt.isWellFormed (cxt : Cxt) : Bool :=
  cxt.env.length == cxt.lvl && cxt.types.length == cxt.lvl

def Cxt.typeAt (cxt : Cxt) (idx : Fin cxt.lvl) : CxtEntry :=
  cxt.types.get (castFin cxt.typesWf.symm idx)

def Cxt.valueAt (cxt : Cxt) (idx : Fin cxt.lvl) : Val :=
  cxt.env.get (castFin cxt.envWf.symm idx)

def Cxt.lookup (cxt : Cxt) (x : Name) : Option (Fin cxt.lvl × Val) :=
  match lookupGo x cxt.types with
  | some (idx, a) =>
      some (castFin cxt.typesWf idx, a)
  | none =>
      none

end BidirTT