import BidirTT

open BidirTT

inductive Expectation where
  | okTy : Tm → Expectation
  | okTyNorm : Tm → Tm → Expectation
  | errContains : String → Expectation

private def renderType (ty : Val) : Except String Tm :=
  BidirTT.quote 0 ty

private def renderNormal (tm : Tm) : Except String Tm := do
  let v ← eval [] tm
  quote 0 v

private def containsText (haystack needle : String) : Bool :=
  needle.isEmpty || (haystack.splitOn needle).length > 1

structure TestCase where
  name   : String
  input  : Raw
  expect : Expectation

def cases : List TestCase := [
  ⟨"U0 is typed by U1",          Examples.univ0,
    .okTy (.univ 1)⟩,
  ⟨"U0 subsumes into U2",        .ann (.univ 0) (.univ 2),
    .okTy (.univ 2)⟩,
  ⟨"Nat is typed by U0",         .nat,
    .okTy (.univ 0)⟩,
  ⟨"succ zero infers Nat",       (.succ .zero),
    .okTy .nat⟩,
  ⟨"id typechecks",             Examples.idAnn,
    .okTy (.pi (.univ 0) (.pi (.var 0) (.var 1)))⟩,
  ⟨"Pi subsumption is contravariant in the domain", .ann
      (.ann
        (.lam "A" (.var "A"))
        (.pi "A" (.univ 1) (.univ 1)))
      (.pi "A" (.univ 0) (.univ 2)),
    .okTy (.pi (.univ 0) (.univ 2))⟩,
  ⟨"const typechecks",          Examples.constAnn,
    .okTy (.pi (.univ 0) (.pi (.univ 0) (.pi (.var 1) (.pi (.var 1) (.var 3)))))⟩,
  ⟨"swap typechecks",           Examples.swapAnn,
    .okTy (.pi (.univ 0) (.pi (.univ 0) (.pi (.sig (.var 1) (.var 1)) (.sig (.var 1) (.var 3)))))⟩,
  ⟨"dependent pair typechecks", Examples.depPairAnn,
    .okTy (.sig (.univ 2) (.var 0))⟩,
  ⟨"dependent pair subsumes into a lifted Sigma", .ann Examples.depPairAnn
      (.sig "A" (.univ 3) (.var "A")),
    .okTy (.sig (.univ 3) (.var 0))⟩,
  ⟨"natElim computes on numerals", Examples.natFoldId,
    .okTyNorm .nat (.succ (.succ .zero))⟩,
  ⟨"unitElim computes on tt",    Examples.unitToNat,
    .okTyNorm .nat (.succ (.succ .zero))⟩,
  ⟨"emptyElim builds absurd maps", Examples.absurdNat,
    .okTy (.pi .empty .nat)⟩,
  ⟨"refl inhabits reflexive identity", Examples.reflZero,
    .okTyNorm (.id .nat .zero .zero) .refl⟩,
  ⟨"idElim computes on refl",    Examples.idElimNat,
    .okTyNorm .nat .zero⟩,
  ⟨"fst infers the first projection", Examples.fstDepPair,
    .okTy (.univ 2)⟩,
  ⟨"snd infers the dependent second projection", Examples.sndDepPair,
    .okTy (.univ 1)⟩,
  ⟨"let infers through definitions", Examples.letUniverse,
    .okTy (.univ 1)⟩,
  ⟨"bad succ rejected", Examples.badSucc,
    .errContains "type mismatch: expected Nat, got U1"⟩,
  ⟨"self application rejected", Examples.omegaAnn,
    .errContains "expected Pi type in application"⟩,
  ⟨"unknown variable rejected", Examples.unknownVar,
    .errContains "unknown variable nope"⟩,
  ⟨"pair mismatch rejected at the Sigma body", Examples.pairMismatch,
    .errContains "type mismatch: expected U1, got U2"⟩,
  ⟨"bad fst rejected", Examples.badFst,
    .errContains "expected Sigma type in .1, got U1"⟩,
  ⟨"bad refl rejected", Examples.badRefl,
    .errContains "refl cannot inhabit"⟩
]

def runCase (tc : TestCase) : IO Bool := do
  match tc.expect, checkTop tc.input with
  | .okTy expectedTy, .ok (_, ty) =>
      match renderType ty with
      | Except.ok actualTy =>
          if actualTy == expectedTy then
            IO.println s!"PASS  {tc.name}"
            pure true
          else
            IO.println s!"FAIL  {tc.name} (expected type {BidirTT.prettyTm expectedTy}, got {BidirTT.prettyTm actualTy})"
            pure false
      | Except.error err =>
          IO.println s!"FAIL  {tc.name} (could not quote type: {err})"
          pure false
  | .okTyNorm expectedTy expectedNf, .ok (tm, ty) =>
      match renderType ty, renderNormal tm with
      | Except.ok actualTy, Except.ok actualNf =>
          if actualTy == expectedTy && actualNf == expectedNf then
            IO.println s!"PASS  {tc.name}"
            pure true
          else if actualTy != expectedTy then
            IO.println s!"FAIL  {tc.name} (expected type {BidirTT.prettyTm expectedTy}, got {BidirTT.prettyTm actualTy})"
            pure false
          else
            IO.println s!"FAIL  {tc.name} (expected nf {BidirTT.prettyTm expectedNf}, got {BidirTT.prettyTm actualNf})"
            pure false
      | Except.error err, _ =>
          IO.println s!"FAIL  {tc.name} (could not quote type: {err})"
          pure false
      | _, Except.error err =>
          IO.println s!"FAIL  {tc.name} (could not normalize term: {err})"
          pure false
  | .okTy expectedTy, .error err =>
      IO.println s!"FAIL  {tc.name} (expected type {BidirTT.prettyTm expectedTy}, got error {err})"
      pure false
  | .okTyNorm expectedTy _, .error err =>
      IO.println s!"FAIL  {tc.name} (expected type {BidirTT.prettyTm expectedTy}, got error {err})"
      pure false
  | .errContains needle, .error err =>
      if containsText err needle then
        IO.println s!"PASS  {tc.name}"
        pure true
      else
        IO.println s!"FAIL  {tc.name} (expected error containing {needle}, got {err})"
        pure false
  | .errContains needle, .ok (_, ty) =>
      match renderType ty with
      | Except.ok actualTy =>
          IO.println s!"FAIL  {tc.name} (expected error containing {needle}, got type {BidirTT.prettyTm actualTy})"
          pure false
      | Except.error err =>
          IO.println s!"FAIL  {tc.name} (expected error containing {needle}, got quote failure {err})"
          pure false

def runInternalSafetyChecks : IO Bool := do
  let malformedEvalOk :=
    match eval [] (.var 0) with
    | Except.error err => containsText err "bad de Bruijn index 0"
    | Except.ok _ => false
  let malformedQuoteOk :=
    match quote 0 (.var 0) with
    | Except.error err => containsText err "bad level 0"
    | Except.ok _ => false
  if malformedEvalOk && malformedQuoteOk then
    IO.println "PASS  malformed core terms are rejected safely"
    pure true
  else
    IO.println "FAIL  malformed core terms are rejected safely"
    pure false

def runPrettyPrinterChecks : IO Bool := do
  match checkTop Examples.idAnn with
  | .ok (tm, ty) =>
      match renderType ty with
      | Except.ok qty =>
          let termText := BidirTT.prettyTm tm
          let typeText := BidirTT.prettyTm qty
          let ok :=
            containsText termText "fun x0 =>" &&
            containsText typeText "Pi (x0 : U0)" &&
            !containsText typeText "#"
          if ok then
            IO.println "PASS  pretty printer rehydrates binder names"
            pure true
          else
            IO.println s!"FAIL  pretty printer rehydrates binder names (term: {termText}, type: {typeText})"
            pure false
      | Except.error err =>
          IO.println s!"FAIL  pretty printer rehydrates binder names (could not quote type: {err})"
          pure false
  | .error err =>
      IO.println s!"FAIL  pretty printer rehydrates binder names (could not elaborate fixture: {err})"
      pure false

def main : IO UInt32 := do
  let results ← cases.mapM runCase
  let safetyOk ← runInternalSafetyChecks
  let prettyOk ← runPrettyPrinterChecks
  let allResults := results ++ [safetyOk, prettyOk]
  let failed := allResults.countP (· == false)
  if failed == 0 then
    IO.println s!"\n{allResults.length} passed, 0 failed"
    pure 0
  else
    IO.println s!"\n{allResults.length - failed} passed, {failed} failed"
    pure 1