import BidirTT

open BidirTT

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

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

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)⟩,
  ⟨"id typechecks",             Examples.idAnn,
    .okTy (.pi (.univ 0) (.pi (.var 0) (.var 1)))⟩,
  ⟨"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))⟩,
  ⟨"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)⟩,
  ⟨"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"⟩
]

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
  | .okTy 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 main : IO UInt32 := do
  let results ← cases.mapM runCase
  let safetyOk ← runInternalSafetyChecks
  let allResults := results ++ [safetyOk]
  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