/-
Conformance test suite generator.

Walks the `<conformance-root>/test-cases/uplc/evaluation/` directory tree and
emits a mirror tree of Lean test files under the output directory (default
`Tests/Conformance/Generated`).  Each leaf test case becomes one `.lean` file;
each branch directory becomes a barrel `.lean` file that imports its children.

Four leaf shapes are emitted, chosen from the first line of `*.uplc.expected`:
  * `parse error`        → `#conformance_import_uplc` with a `Parsing error` guard;
  * `evaluation failure` → `#import_uplc` + `programEvalsToError`, except for the
                            source paths listed in `scopeErrorPaths`, which get a
                            `#conformance_import_uplc` + `Parsing error` guard
                            (see that declaration for why);
  * otherwise            → `#import_uplc` (program + expected) + `programsEvalEquiv`
                            + `budgetMatches` (cpu/mem read from `*.uplc.budget.expected`).

Usage:
  gen_conformance_tests <conformance-root> [--out <dir>] [--embed-root <prefix>]
                                           [--exclude-not-implemented]
-/

namespace GenConformance

abbrev Path := System.FilePath

structure Config where
  conformanceRoot : Path
  embedRoot       : String
  outRoot         : Path
  namespacePrefix : List String
  /-- Source-relative path prefixes (each a list of directory segments under
      `test-cases/uplc/evaluation/`) to skip when generating. -/
  excludedPaths   : List (List String) := []

/-- Source-relative path prefixes for test categories whose corresponding
    features are not yet implemented in this Lean formalization. Selected by
    `--exclude-not-implemented`. -/
def notImplementedPaths : List (List String) := [
  ["builtin", "constant", "array"],
  ["builtin", "constant", "value"],
  ["builtin", "semantics", "indexArray"],
  ["builtin", "semantics", "insertCoin"],
  ["builtin", "semantics", "lengthOfArray"],
  ["builtin", "semantics", "listToArray"],
  ["builtin", "semantics", "lookupCoin"],
  ["builtin", "semantics", "scaleValue"],
  ["builtin", "semantics", "unionValue"],
  ["builtin", "semantics", "unValueData"],
  ["builtin", "semantics", "valueContains"],
  ["builtin", "semantics", "valueData"]
]

/-- Source-relative path prefixes whose expected `evaluation failure` is caused by a free (unbound) variable.

    The reference implementation parses UPLC into a named representation and only
    fails when converting it to de Bruijn form, so it reports these as evaluation
    failures. -/
def scopeErrorPaths : List (List String) := [
  ["term", "var"]
]

def header : String :=
  "-- Auto-generated by GenConformanceTests.lean — do not edit manually."

private def lower (s : String) : String := s.map Char.toLower

private def capFirst (s : String) : String :=
  if s.isEmpty then s else s.set 0 (s.get 0).toUpper

/-- Replace characters that aren't valid in Lean identifiers (currently `-` and `+`)
    with underscores. -/
private def sanitize (s : String) : String :=
  s.map fun c => if c == '-' || c == '+' then '_' else c

/-- Source directory segment → Lean module / file segment. -/
def toLeanName (s : String) : String := capFirst (sanitize s)

/-- Source directory segment → Lean identifier component. -/
def toIdent (s : String) : String := sanitize s

/-- Whether a source directory or `.uplc` file name may go into the generated
    files: a letter followed by letters, digits, `_`, `-` and `+`. The names
    become Lean identifiers, module names, doc comments and string literals, and
    they come from the plutus checkout, so anything else is refused rather than
    escaped. -/
def isSafeName (s : String) : Bool :=
  match s.toList with
  | c :: cs => c.isAlpha && cs.all fun c => c.isAlphanum || c == '_' || c == '-' || c == '+'
  | []      => false

/-! ## Parsing the expected/budget files -/

inductive ExpectedKind where
  | parseError
  | evalFailure
  | scopeError
  | success (cpu mem : Nat)

private def firstLine (s : String) : String :=
  match s.splitOn "\n" with
  | l :: _ => l.trim
  | []     => ""

/-- Parse `key: NN` from a budget-file line, tolerating leading `(`, `{`, `|`, ws. -/
private def parseBudgetField (key : String) (line : String) : Option Nat :=
  let line := line.dropWhile fun c =>
    c == '(' || c == '{' || c == '|' || c == ' ' || c == '\t'
  if line.startsWith key then
    let rest := (line.drop key.length).trim
    let digits := rest.takeWhile (·.isDigit)
    digits.toNat?
  else none

private def parseBudget (s : String) : Option (Nat × Nat) :=
  let lines := s.splitOn "\n"
  match lines.findSome? (parseBudgetField "cpu:"),
        lines.findSome? (parseBudgetField "mem:") with
  | some c, some m => some (c, m)
  | _, _           => none

private def classifyExpected (srcPathRel : List String) (expectedContent budgetContent : String) : Except String ExpectedKind :=
  let line := firstLine expectedContent
  if line.startsWith "parse error" then
    .ok .parseError
  else if line.startsWith "evaluation failure" then
    if scopeErrorPaths.any (·.isPrefixOf srcPathRel) then .ok .scopeError
    else .ok .evalFailure
  else
    match parseBudget budgetContent with
    | some (c, m) => .ok (.success c m)
    | none        => .error "could not parse budget file"

/-! ## Path helpers -/

private def withLeanExt (p : Path) : Path := ⟨p.toString ++ ".lean"⟩

/-- File path of the leaf or barrel `.lean` for module path `mp`.
    For `mp = []`, this is the outer barrel one level above `outRoot`. -/
private def leanFileFor (cfg : Config) (mp : List String) : Path :=
  match mp with
  | [] =>
      let parent := cfg.outRoot.parent.getD ⟨"."⟩
      let name := cfg.outRoot.fileName.getD "Generated"
      parent / (name ++ ".lean")
  | _ =>
      withLeanExt (mp.foldl (fun (p : Path) (s : String) => p / s) cfg.outRoot)

/-- Directory containing the children of module path `mp` (a branch dir). -/
private def childrenDirFor (cfg : Config) (mp : List String) : Path :=
  mp.foldl (fun (p : Path) (s : String) => p / s) cfg.outRoot

/-! ## Templates -/

private def fullModuleName (cfg : Config) (mp : List String) : String :=
  ".".intercalate (cfg.namespacePrefix ++ mp)

private def fullIdent (srcPathRel : List String) : String :=
  "_".intercalate (srcPathRel.map toIdent)

private def leafEmbedPath
    (cfg : Config) (srcPathRel : List String) (stem suffix : String) : String :=
  let dir := "/".intercalate srcPathRel
  s!"{cfg.embedRoot}/test-cases/uplc/evaluation/{dir}/{stem}{suffix}"

private def leafSuccessContent
    (cfg : Config) (srcPathRel : List String) (mp : List String)
    (stem : String) (cpu mem : Nat) : String :=
  let ns := fullModuleName cfg mp
  let id := fullIdent srcPathRel
  let p1 := leafEmbedPath cfg srcPathRel stem ".uplc"
  let p2 := leafEmbedPath cfg srcPathRel stem ".uplc.expected"
  s!"\
{header}
import Tests.Conformance.ConformanceUtils
import PlutusCore.UPLC.ScriptEncoding.Basic

namespace {ns}

open Tests.Conformance

/-- info: Successfully decoded textual '{p1}' -/
#guard_msgs in
#import_uplc {id} PlutusV3 textual \"{p1}\"

/-- info: Successfully decoded textual '{p2}' -/
#guard_msgs in
#import_uplc {id}_expected PlutusV3 textual \"{p2}\"

/-- info: true -/
#guard_msgs in
#eval programsEvalEquiv {id} {id}_expected

/-- info: true -/
#guard_msgs in
#eval! budgetMatches {id} {cpu} {mem}

end {ns}
"

private def leafEvalFailureContent
    (cfg : Config) (srcPathRel : List String) (mp : List String)
    (stem : String) : String :=
  let ns := fullModuleName cfg mp
  let id := fullIdent srcPathRel
  let p1 := leafEmbedPath cfg srcPathRel stem ".uplc"
  s!"\
{header}
import Tests.Conformance.ConformanceUtils
import PlutusCore.UPLC.ScriptEncoding.Basic

namespace {ns}

open Tests.Conformance

/-- info: Successfully decoded textual '{p1}' -/
#guard_msgs in
#import_uplc {id} PlutusV3 textual \"{p1}\"

/-- info: true -/
#guard_msgs in
#eval programEvalsToError {id}

end {ns}
"

/-- Leaf for an `evaluation failure` that this formalization rejects at decoding
    time rather than at evaluation time; see `scopeErrorPaths`. -/
private def leafScopeErrorContent
    (cfg : Config) (srcPathRel : List String) (mp : List String)
    (stem : String) : String :=
  let ns := fullModuleName cfg mp
  let id := fullIdent srcPathRel
  let p1 := leafEmbedPath cfg srcPathRel stem ".uplc"
  s!"\
{header}
import Tests.Conformance.ConformanceUtils
import PlutusCore.UPLC.ScriptEncoding.Basic

namespace {ns}

open Tests.Conformance

/-- error: Parsing error -/
#guard_msgs in
#conformance_import_uplc {id} \"{p1}\"

end {ns}
"

private def leafParseErrorContent
    (cfg : Config) (srcPathRel : List String) (mp : List String)
    (stem : String) : String :=
  let ns := fullModuleName cfg mp
  let id := fullIdent srcPathRel
  let p1 := leafEmbedPath cfg srcPathRel stem ".uplc"
  s!"\
{header}
import Tests.Conformance.ConformanceUtils
import PlutusCore.UPLC.ScriptEncoding.Basic

namespace {ns}

open Tests.Conformance

/-- error: Parsing error -/
#guard_msgs in
#conformance_import_uplc {id} \"{p1}\"

end {ns}
"

private def barrelContent (cfg : Config) (mp : List String) (children : List String) : String :=
  let ns := fullModuleName cfg mp
  if children.isEmpty then
    s!"{header}\n\n-- Barrel for {ns}\n"
  else
    let imports := "\n".intercalate (children.map fun c => s!"import {ns}.{c}")
    s!"{header}\n\n{imports}\n\n-- Barrel for {ns}\n"

/-! ## Walker -/

/-- A source path is excluded if it equals or is a descendant of any excluded
    prefix in `cfg.excludedPaths`. -/
private def isExcluded (cfg : Config) (srcPathRel : List String) : Bool :=
  cfg.excludedPaths.any fun ex => ex.isPrefixOf srcPathRel

/-- Remove any previously-generated output for excluded prefixes so the output
    tree remains consistent when re-running with `--exclude-not-implemented`. -/
private def cleanExcludedOutputs (cfg : Config) : IO Unit := do
  for ex in cfg.excludedPaths do
    let leanPath := ex.map toLeanName
    let dir := childrenDirFor cfg leanPath
    let file := leanFileFor cfg leanPath
    if ← dir.isDir then
      IO.FS.removeDirAll dir
    if ← file.pathExists then
      IO.FS.removeFile file

/-- Fails on a name that `isSafeName` refuses. The name is printed quoted, so
    that it cannot inject anything into the log either. -/
private def checkName (srcDir : Path) (name : String) : IO Unit :=
  unless isSafeName name do
    throw <| IO.userError s!"unsupported test name {name.quote} in {srcDir.toString}: \
      names must start with a letter and contain only letters, digits, '_', '-' and '+'"

mutual
  partial def walkBranch (cfg : Config) (srcDir : Path)
      (srcPathRel mp : List String) (subdirs : List String) : IO Unit := do
    let childrenDir := childrenDirFor cfg mp
    IO.FS.createDirAll childrenDir
    let allSorted := (subdirs.toArray.qsort (fun a b => lower a < lower b)).toList
    let sortedSubs := allSorted.filter fun sub => ¬ isExcluded cfg (srcPathRel ++ [sub])
    for sub in sortedSubs do
      checkName srcDir sub
      let childMp := mp ++ [toLeanName sub]
      let childSrc := srcPathRel ++ [sub]
      walk cfg (srcDir / sub) childSrc childMp
    let childModules := sortedSubs.map toLeanName
    IO.FS.writeFile (leanFileFor cfg mp) (barrelContent cfg mp childModules)

  partial def walk (cfg : Config) (srcDir : Path)
      (srcPathRel mp : List String) : IO Unit := do
    if isExcluded cfg srcPathRel then return
    let entries ← srcDir.readDir
    let mut subdirs : Array String := #[]
    let mut uplcStem : Option String := none
    for e in entries do
      let name := e.fileName
      let isDir ← (srcDir / name).isDir
      if isDir then
        subdirs := subdirs.push name
      else if name.endsWith ".uplc"
              && ¬ name.endsWith ".uplc.expected"
              && ¬ name.endsWith ".uplc.budget.expected" then
        uplcStem := some (name.dropRight ".uplc".length)
    match uplcStem with
    | some stem =>
        checkName srcDir stem
        let expectedPath := srcDir / s!"{stem}.uplc.expected"
        let budgetPath   := srcDir / s!"{stem}.uplc.budget.expected"
        let expectedContent ← IO.FS.readFile expectedPath
        let budgetContent ← (try IO.FS.readFile budgetPath catch _ => pure "")
        match classifyExpected srcPathRel expectedContent budgetContent with
        | .ok kind =>
            let outFile := leanFileFor cfg mp
            if let some parent := outFile.parent then
              IO.FS.createDirAll parent
            let content := match kind with
              | .parseError    => leafParseErrorContent cfg srcPathRel mp stem
              | .evalFailure   => leafEvalFailureContent cfg srcPathRel mp stem
              | .scopeError    => leafScopeErrorContent cfg srcPathRel mp stem
              | .success c m   => leafSuccessContent cfg srcPathRel mp stem c m
            IO.FS.writeFile outFile content
        | .error msg =>
            throw <| IO.userError s!"failed to classify {srcDir.toString}: {msg}"
    | none =>
        walkBranch cfg srcDir srcPathRel mp subdirs.toList
end

/-! ## CLI -/

def usage : String :=
"Usage: gen_conformance_tests <conformance-root> [--out <dir>] [--embed-root <prefix>]
                                                 [--exclude-not-implemented]

Walks the conformance test directory tree and emits Lean test files.

Arguments:
  <conformance-root>          Path to the plutus-conformance/ directory
                              (the one containing test-cases/).

Options:
  --out <dir>                 Output directory for generated .lean files.
                              Default: Tests/Conformance/Generated
  --embed-root <prefix>       Path prefix to embed in generated #import_uplc strings.
                              Default: same as <conformance-root>.
  --exclude-not-implemented   Skip generation for test categories whose underlying
                              features are not yet implemented in this Lean
                              formalization (see `notImplementedPaths`).
                              Existing output for those categories is removed.
"

partial def parseArgs : List String →
    Option (List String × Option String × Option String × Bool)
  | [] => some ([], none, none, false)
  | "--out" :: v :: rest => do
      let (pos, _, e, x) ← parseArgs rest
      some (pos, some v, e, x)
  | "--embed-root" :: v :: rest => do
      let (pos, o, _, x) ← parseArgs rest
      some (pos, o, some v, x)
  | "--exclude-not-implemented" :: rest => do
      let (pos, o, e, _) ← parseArgs rest
      some (pos, o, e, true)
  | x :: rest =>
      if x.startsWith "--" then none
      else do
        let (pos, o, e, x') ← parseArgs rest
        some (x :: pos, o, e, x')

def run (args : List String) : IO UInt32 := do
  match parseArgs args with
  | none =>
      IO.eprintln usage
      return 2
  | some ([conf], outOpt, embedOpt, excludeNI) =>
      let conformanceRoot : Path := ⟨conf⟩
      let outRoot : Path :=
        match outOpt with
        | some s => ⟨s⟩
        | none   => ⟨"Tests/Conformance/Generated"⟩
      let embedRoot := embedOpt.getD conf
      let cfg : Config := {
        conformanceRoot
        embedRoot
        outRoot
        namespacePrefix := ["Tests", "Conformance", "Generated"]
        excludedPaths   := if excludeNI then notImplementedPaths else []
      }
      let evaluationDir := conformanceRoot / "test-cases" / "uplc" / "evaluation"
      unless ← evaluationDir.isDir do
        IO.eprintln s!"Error: not a directory: {evaluationDir.toString}"
        return 1
      IO.FS.createDirAll cfg.outRoot
      cleanExcludedOutputs cfg
      walk cfg evaluationDir [] []
      IO.println s!"Wrote generated tests under {cfg.outRoot.toString} \
                    (outer barrel: {(leanFileFor cfg []).toString})."
      return 0
  | _ =>
      IO.eprintln usage
      return 2

end GenConformance

def main (args : List String) : IO UInt32 := GenConformance.run args