import Lean

open Lean IO

--------------------------------------------------------------------------------
-- 1. Data structures
--------------------------------------------------------------------------------

inductive KBCommand where
  | Init        (dir : String)
  | AddType     (name : String) (parent : Option String) (domain : String)
  | AddInstance (name : String) (typeName : String)      (domain : String)
  | AddProp     (name : String) (domType : String) (range : String) (domain : String)
  | AddRel      (name : String) (dom1 : String) (dom2 : String)     (domain : String)
  | AddFact     (fact : String)                                      (domain : String)
  | AddRule     (name : String) (statement : String)                 (domain : String)
  | AddDef      (name : String) (params : String) (type : String) (body : String) (domain : String)
  | Remove      (name : String)
  | List        (what : String) (filter : Option String)
  | ShowDef     (name : String)
  | Get         (expr : String)
  | Search      (query : String)
  | Deps        (name : String)
  | Eval        (expr : String)
  | Query       (prop : String) (tactics : Option String)
  | Check
  | Commit      (msg : String)
  | Status
  | Domains
  | DomainsStats
  | Help
  | Exit
deriving Repr

/-- A single KB declaration.
    Written to disk as a structured metadata comment immediately before the Lean source:
      --@kb kind=<k> name=<n> domain=<d>
      <lean source line(s)>
    This lets the tool reconstruct full state on startup by scanning domain files. -/
structure KBDecl where
  kind   : String   -- "type" | "instance" | "prop" | "rel" | "fact" | "rule" | "def"
  name   : String
  lean   : String   -- Lean 4 source (may be multi-line)
  domain : String   -- stem of the .lean file (without extension)
deriving Repr, Inhabited

structure KBState where
  root  : System.FilePath
  decls : Array KBDecl

/-- Result of a command — drives exit codes in non-interactive use. -/
structure CmdResult where
  ok       : Bool   -- false → exit code 1
  continue : Bool   -- false → user typed 'exit', stop the REPL

def cmdOk   : CmdResult := { ok := true,  continue := true  }
def cmdErr  : CmdResult := { ok := false, continue := true  }
def cmdQuit : CmdResult := { ok := true,  continue := false }

--------------------------------------------------------------------------------
-- 2. Utility functions
--------------------------------------------------------------------------------

def strContains (haystack needle : String) : Bool :=
  if needle.isEmpty then true
  else (haystack.splitOn needle).length > 1

def dedupStrings (xs : List String) : List String :=
  xs.foldl (fun acc x => if acc.contains x then acc else acc ++ [x]) []

def strToLower (s : String) : String :=
  s.map fun c =>
    if c.toNat >= 'A'.toNat && c.toNat <= 'Z'.toNat
    then Char.ofNat (c.toNat + 32) else c

/-- Split on the LAST occurrence of `sep` and return (before, after). -/
def splitOnLast (s sep : String) : Option (String × String) :=
  let segs := s.splitOn sep
  if segs.length < 2 then none
  else some
    ((segs.dropLast |> String.intercalate sep).trim,
     segs.getLast!.trim)

--------------------------------------------------------------------------------
-- 3. Metadata format and persistence
--
-- Every declaration on disk is stored as:
--   --@kb kind=<k> name=<n> domain=<d>
--   <lean source line 1>
--   [<lean source line 2> ...]
--   (blank line)
--
-- Parsing all domain files on startup fully reconstructs KBState.decls.
--------------------------------------------------------------------------------

def metaPrefix : String := "--@kb "

def encodeMetaLine (d : KBDecl) : String :=
  s!"{metaPrefix}kind={d.kind} name={d.name} domain={d.domain}"

/-- Parse a "--@kb ..." header. Returns (kind, name, domain) or none. -/
def parseMetaLine (line : String) : Option (String × String × String) := do
  guard (line.startsWith metaPrefix)
  let fields := (line.drop metaPrefix.length).splitOn " "
  let mut kind := ""; let mut name := ""; let mut domain := ""
  for f in fields do
    if f.startsWith "kind="   then kind   := f.drop "kind=".length
    if f.startsWith "name="   then name   := f.drop "name=".length
    if f.startsWith "domain=" then domain := f.drop "domain=".length
  if kind.isEmpty || name.isEmpty || domain.isEmpty then none
  else some (kind, name, domain)

/-- Extract all --@kb-annotated declarations from the text of one domain file. -/
def parseDomainFile (contents : String) : Array KBDecl :=
  let lines := contents.splitOn "\n"
  let mut decls : Array KBDecl := #[]
  let mut i := 0
  while i < lines.length do
    match parseMetaLine lines[i]! with
    | none => i := i + 1
    | some (kind, name, domain) =>
        i := i + 1
        let mut leanLines : List String := []
        while i < lines.length && !(lines[i]!.startsWith metaPrefix) do
          leanLines := leanLines ++ [lines[i]!]
          i := i + 1
        let lean :=
          (leanLines.reverse.dropWhile String.isEmpty).reverse
          |> String.intercalate "\n"
        decls := decls.push { kind, name, lean, domain }
  decls

/-- Scan all domain .lean files in `root` and reconstruct the declaration list. -/
def loadState (root : System.FilePath) : IO (Array KBDecl) := do
  let reserved := ["KnowledgeBase", "lakefile", "_kb_tmp"]
  let entries  ← System.FilePath.readDir root
  let stems := entries.toList.filterMap fun e =>
    let fname := e.fileName
    if fname.endsWith ".lean" then
      let stem := fname.dropRight ".lean".length
      if reserved.contains stem then none else some stem
    else none
  let mut all : Array KBDecl := #[]
  for stem in dedupStrings stems do
    let path := root / (stem ++ ".lean")
    if ← path.pathExists then
      all := all ++ parseDomainFile (← IO.FS.readFile path)
  return all

--------------------------------------------------------------------------------
-- 4. Command parsing
--------------------------------------------------------------------------------

def parseCommand (line : String) : Option KBCommand :=
  let normalised := line.trim.replace "→" "->"
  let parts := normalised.splitOn " " |>.filter (not ∘ String.isEmpty)
  let raw   := line.trim
  match parts with
  | ["init", dir] => some (.Init dir)
  -- add type
  | ["add", "type", name, "in", dom] =>
      some (.AddType name none dom)
  | ["add", "type", name, "extends", par, "in", dom] =>
      some (.AddType name (some par) dom)
  -- add instance
  | ["add", "instance", name, ":", ty, "in", dom] =>
      some (.AddInstance name ty dom)
  -- add prop
  | ["add", "prop", name, ":", dt, "->", rng, "in", dom] =>
      some (.AddProp name dt rng dom)
  -- add rel
  | ["add", "rel", name, ":", d1, "->", d2, "->", "Prop", "in", dom] =>
      some (.AddRel name d1 d2 dom)
  -- add fact  (suffix style — last " in " is the domain separator)
  | "add" :: "fact" :: _ =>
      match splitOnLast (raw.drop "add fact ".length) " in " with
      | some (fact, dom) => some (.AddFact fact dom)
      | none             => none
  -- add rule
  | "add" :: "rule" :: name :: ":" :: _ =>
      let afterColon := raw.drop ("add rule " ++ name ++ " : ").length
      match splitOnLast afterColon " in " with
      | some (stmt, dom) => some (.AddRule name stmt dom)
      | none             => none
  -- add def
  | "add" :: "def" :: name :: _ =>
      let afterName := raw.drop ("add def " ++ name ++ " ").length
      match splitOnLast afterName " in " with
      | none => none
      | some (defBody, dom) =>
          match splitOnLast defBody " := " with
          | none => none
          | some (lhs, body) =>
              match splitOnLast lhs " : " with
              | none => none
              | some (params, ty) => some (.AddDef name params ty body dom)
  -- rm
  | ["rm", name]            => some (.Remove name)
  -- list
  | ["list", what]          => some (.List what none)
  | ["list", what, "of", f] => some (.List what (some f))
  -- show / get / search / deps
  | ["show", "def", name]   => some (.ShowDef name)
  | "get" :: _              => some (.Get    (raw.drop "get ".length))
  | "search" :: _           => some (.Search (raw.drop "search ".length))
  | ["deps", name]          => some (.Deps name)
  -- eval
  | "eval" :: _             => some (.Eval (raw.drop "eval ".length))
  -- query [using <tactics>]  — split on last " using "
  | "query" :: _ =>
      let body := raw.drop "query ".length
      match splitOnLast body " using " with
      | some (prop, tacs) => some (.Query prop (some tacs))
      | none              => some (.Query body none)
  -- build / git
  | ["check"]               => some .Check
  | "commit" :: _           => some (.Commit (raw.drop "commit ".length))
  | ["status"]              => some .Status
  -- domain info
  | ["domains"]             => some .Domains
  | ["domains", "stats"]    => some .DomainsStats
  -- meta
  | ["help"]                => some .Help
  | ["exit"]                => some .Exit
  | _                       => none

--------------------------------------------------------------------------------
-- 5. File-system helpers
--------------------------------------------------------------------------------

def domainPath (root : System.FilePath) (domain : String) : System.FilePath :=
  root / (domain ++ ".lean")

/-- Write a declaration to disk: metadata header then Lean source, then blank line. -/
def appendToDomain (root : System.FilePath) (d : KBDecl) : IO Unit := do
  let path := domainPath root d.domain
  unless (← path.pathExists) do
    IO.FS.writeFile path s!"-- Knowledge Base: {d.domain}\n\n"
  let h ← IO.FS.Handle.mk path .append
  h.putStrLn (encodeMetaLine d)
  h.putStrLn d.lean
  h.putStrLn ""
  h.flush

def runProc (cmd : String) (args : Array String) (cwd : Option String := none)
    : IO (UInt32 × String) := do
  let child ← IO.Process.spawn {
    cmd    := cmd; args := args; cwd := cwd
    stdout := .piped; stderr := .piped
  }
  let stdoutTask ← IO.asTask child.stdout.readToEnd Task.Priority.dedicated
  let stderrOut  ← child.stderr.readToEnd
  let stdoutOut  ← IO.ofExcept stdoutTask.get
  let code       ← child.wait
  return (code, stdoutOut ++ stderrOut)

def runProcPrint (cmd : String) (args : Array String) (cwd : Option String := none)
    : IO UInt32 := do
  let (code, out) ← runProc cmd args cwd
  unless out.isEmpty do println out
  return code

--------------------------------------------------------------------------------
-- 6. Lean snippet generators
--------------------------------------------------------------------------------

def genType (name : String) (parent : Option String) : String :=
  match parent with
  | none   => s!"opaque {name} : Type"
  | some p =>
      s!"opaque {name} : Type\n" ++
      s!"-- {name} extends {p}\n" ++
      s!"axiom {name}.val : {name} → {p}"

def genInstance (name typeName : String) : String :=
  s!"opaque {name} : {typeName}"

def genProp (name domType range : String) : String :=
  s!"opaque {name} : {domType} → {range}"

def genRel (name dom1 dom2 : String) : String :=
  s!"opaque {name} : {dom1} → {dom2} → Prop"

def genFact (fact : String) : String :=
  let byParen := fact.splitOn "("
  if byParen.length < 2 then
    let byEq := fact.splitOn "="
    if byEq.length >= 2 then
      let lhs := (byEq.headD "").trim
      let rhs := (byEq.tail.headD "").trim
      s!"axiom fact_{lhs} : {lhs} = {rhs}"
    else s!"-- unrecognised fact: {fact}"
  else
    let funcName   := (byParen.headD "").trim
    let afterOpen  := byParen.tail |> String.intercalate "("
    let byClose    := afterOpen.splitOn ")"
    let argsStr    := (byClose.headD "").trim
    let args       := argsStr.splitOn "," |>.map String.trim
    let tag        := args.foldl (· ++ "_" ++ ·) funcName
    let argStr     := String.intercalate " " args
    let afterClose := byClose.tail |> String.intercalate ")"
    let byEq       := afterClose.splitOn "="
    if byEq.length >= 2 then
      s!"axiom fact_{tag} : {funcName} {argStr} = {(byEq.tail.headD "").trim}"
    else
      s!"axiom fact_{tag} : {funcName} {argStr}"

def genRule (name stmt : String) : String := s!"axiom {name} : {stmt}"
def genDef  (name params type body : String) : String :=
  s!"def {name} ({params}) : {type} := {body}"

--------------------------------------------------------------------------------
-- 7. Duplicate check
--------------------------------------------------------------------------------

def isDuplicate (state : KBState) (kind name : String) : Bool :=
  state.decls.any (fun d => d.kind == kind && d.name == name)

--------------------------------------------------------------------------------
-- 8. Add handlers
--------------------------------------------------------------------------------

def addType (stateRef : IO.Ref KBState) (name : String) (parent : Option String)
    (domain : String) : IO CmdResult := do
  let state ← stateRef.get
  if isDuplicate state "type" name then
    println s!"Error: type '{name}' already exists."; return cmdErr
  let d : KBDecl := { kind := "type", name, lean := genType name parent, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added type '{name}'{parent.map (" (extends " ++ · ++ ")") |>.getD ""}  →  {domain}.lean"
  return cmdOk

def addInstance (stateRef : IO.Ref KBState) (name typeName domain : String) : IO CmdResult := do
  let state ← stateRef.get
  if isDuplicate state "instance" name then
    println s!"Error: instance '{name}' already exists."; return cmdErr
  let d : KBDecl := { kind := "instance", name, lean := genInstance name typeName, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added instance '{name} : {typeName}'  →  {domain}.lean"
  return cmdOk

def addProp (stateRef : IO.Ref KBState) (name domType range domain : String) : IO CmdResult := do
  let state ← stateRef.get
  if isDuplicate state "prop" name then
    println s!"Error: property '{name}' already exists."; return cmdErr
  let d : KBDecl := { kind := "prop", name, lean := genProp name domType range, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added property '{name} : {domType} → {range}'  →  {domain}.lean"
  return cmdOk

def addRel (stateRef : IO.Ref KBState) (name dom1 dom2 domain : String) : IO CmdResult := do
  let state ← stateRef.get
  if isDuplicate state "rel" name then
    println s!"Error: relation '{name}' already exists."; return cmdErr
  let d : KBDecl := { kind := "rel", name, lean := genRel name dom1 dom2, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added relation '{name} : {dom1} → {dom2} → Prop'  →  {domain}.lean"
  return cmdOk

def addFact (stateRef : IO.Ref KBState) (fact domain : String) : IO CmdResult := do
  let state ← stateRef.get
  let lean := genFact fact
  -- The canonical name is the axiom name embedded in the generated source.
  let axiomName :=
    if lean.startsWith "axiom " then
      (lean.drop "axiom ".length).splitOn " " |>.headD ("fact:" ++ fact)
    else "fact:" ++ fact
  if isDuplicate state "fact" axiomName then
    println s!"Warning: fact '{axiomName}' already exists."; return cmdErr
  let d : KBDecl := { kind := "fact", name := axiomName, lean, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added fact: {fact}  →  {domain}.lean"
  return cmdOk

def addRule (stateRef : IO.Ref KBState) (name stmt domain : String) : IO CmdResult := do
  let state ← stateRef.get
  if isDuplicate state "rule" name then
    println s!"Error: rule '{name}' already exists."; return cmdErr
  let d : KBDecl := { kind := "rule", name, lean := genRule name stmt, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added rule '{name}'  →  {domain}.lean"
  return cmdOk

def addDef (stateRef : IO.Ref KBState) (name params type body domain : String) : IO CmdResult := do
  let state ← stateRef.get
  if isDuplicate state "def" name then
    println s!"Error: definition '{name}' already exists. Use 'rm {name}' first."; return cmdErr
  let d : KBDecl := { kind := "def", name, lean := genDef name params type body, domain }
  appendToDomain state.root d
  stateRef.modify fun s => { s with decls := s.decls.push d }
  println s!"Added definition '{name}'  →  {domain}.lean"
  return cmdOk

--------------------------------------------------------------------------------
-- 9. Remove — metadata-block aware
--
-- Rewrites the domain file, skipping the entire --@kb block for the named decl.
-- This is atomic and correct regardless of how many lines the Lean source spans.
--------------------------------------------------------------------------------

def rewriteFileRemovingDecl (path : System.FilePath) (target : String) : IO Unit := do
  if !(← path.pathExists) then return
  let lines := (← IO.FS.readFile path).splitOn "\n"
  let mut kept : List String := []
  let mut i := 0
  while i < lines.length do
    match parseMetaLine lines[i]! with
    | some (_, mname, _) =>
        if mname == target then
          -- Skip this metadata line and all following non-metadata content.
          i := i + 1
          while i < lines.length && !(lines[i]!.startsWith metaPrefix) do
            i := i + 1
          -- Skip exactly one trailing blank line if present.
          if i < lines.length && lines[i]!.isEmpty then i := i + 1
        else
          kept := kept ++ [lines[i]!]
          i := i + 1
    | none =>
        kept := kept ++ [lines[i]!]
        i := i + 1
  IO.FS.writeFile path (String.intercalate "\n" kept)

def removeDecl (stateRef : IO.Ref KBState) (name : String) : IO CmdResult := do
  let state ← stateRef.get
  let targets := state.decls.filter (fun d => d.name == name)
  if targets.isEmpty then
    println s!"Error: '{name}' not found."; return cmdErr
  let dependents := state.decls.filter (fun d =>
    d.name != name && strContains d.lean name)
  unless dependents.isEmpty do
    println s!"Warning: the following declarations reference '{name}':"
    for dep in dependents do println s!"  [{dep.kind}] {dep.name}  ({dep.domain})"
    println "  Consider removing dependents first."
  for domain in dedupStrings (targets.map (·.domain)).toList do
    let _ ← rewriteFileRemovingDecl (domainPath state.root domain) name
  stateRef.modify fun s => { s with decls := s.decls.filter (fun d => d.name != name) }
  println s!"Removed '{name}'."
  return cmdOk

--------------------------------------------------------------------------------
-- 10. List / ShowDef / Search / Deps / Get
--------------------------------------------------------------------------------

def normKind (w : String) : String :=
  match w with
  | "types" => "type" | "instances" => "instance" | "props" => "prop"
  | "rels"  => "rel"  | "facts"     => "fact"      | "rules" => "rule"
  | "defs"  => "def"  | _           => w

def listDecls (stateRef : IO.Ref KBState) (what : String) (filter : Option String)
    : IO CmdResult := do
  let state ← stateRef.get
  let nk := normKind what
  let matchKind   : KBDecl → Bool :=
    if nk == "all" then fun _ => true else fun d => d.kind == nk
  let matchFilter : KBDecl → Bool :=
    match filter with
    | none   => fun _ => true
    | some f =>
        let fl := strToLower f
        fun d => strContains (strToLower d.name)   fl ||
                 strContains (strToLower d.domain) fl ||
                 strContains (strToLower d.lean)   fl
  let hits := state.decls.filter (fun d => matchKind d && matchFilter d)
  if hits.isEmpty then println s!"(no {what} found)"
  else for d in hits do println s!"[{d.kind}]  {d.name}   ({d.domain})"
  return cmdOk

def showDef (stateRef : IO.Ref KBState) (name : String) : IO CmdResult := do
  let state ← stateRef.get
  let hits := state.decls.filter (fun d => d.name == name)
  if hits.isEmpty then println s!"Error: '{name}' not found."; return cmdErr
  for d in hits do
    println s!"-- [{d.kind}] {d.name}  (domain: {d.domain})"
    println d.lean
  return cmdOk

def searchDecls (stateRef : IO.Ref KBState) (query : String) : IO CmdResult := do
  let state ← stateRef.get
  let q := strToLower query
  let hits := state.decls.filter fun d =>
    strContains (strToLower d.name)   q ||
    strContains (strToLower d.lean)   q ||
    strContains (strToLower d.domain) q
  if hits.isEmpty then
    println s!"No results for '{query}'."
  else do
    println s!"{hits.size} result(s) for '{query}':"
    for d in hits do
      println s!"  [{d.kind}] {d.name}  ({d.domain})"
      for l in d.lean.splitOn "\n" do
        if strContains (strToLower l) q then println s!"    {l.trim}"
  return cmdOk

def showDeps (stateRef : IO.Ref KBState) (name : String) : IO CmdResult := do
  let state ← stateRef.get
  match state.decls.find? (fun d => d.name == name) with
  | none => println s!"Error: '{name}' not found."; return cmdErr
  | some d =>
    let forward := state.decls.filter fun o =>
      o.name != name && strContains d.lean o.name
    let reverse := state.decls.filter fun o =>
      o.name != name && strContains o.lean name
    println s!"Dependencies of [{d.kind}] {name}  ({d.domain}):"
    if forward.isEmpty then println "  uses:     (none)"
    else do
      println "  uses:"
      for dep in forward do println s!"    [{dep.kind}] {dep.name}  ({dep.domain})"
    if reverse.isEmpty then println "  used by:  (none)"
    else do
      println "  used by:"
      for dep in reverse do println s!"    [{dep.kind}] {dep.name}  ({dep.domain})"
    return cmdOk

--------------------------------------------------------------------------------
-- 11. Eval / Get / Query
--------------------------------------------------------------------------------

def buildImports (state : KBState) : IO String := do
  let domains := dedupStrings (state.decls.map (·.domain)).toList
  let mut imports := ""
  for domain in domains do
    if ← (domainPath state.root domain).pathExists then
      imports := imports ++ s!"import {domain}\n"
  return imports

def runLeanSnippet (state : KBState) (snippet : String) : IO (UInt32 × String) := do
  let imports ← buildImports state
  let tmp := state.root / "_kb_tmp.lean"
  IO.FS.writeFile tmp (imports ++ "\n" ++ snippet ++ "\n")
  let result ← runProc "lean" #[tmp.toString] (some state.root.toString)
  try IO.FS.removeFile tmp catch _ => pure ()
  return result

def evalExpr (stateRef : IO.Ref KBState) (expr : String) : IO CmdResult := do
  let state ← stateRef.get
  let (code, out) ← runLeanSnippet state s!"#eval {expr}"
  if code == 0 then println out.trim; return cmdOk
  else do println s!"Lean error evaluating '{expr}':\n{out.trim}"; return cmdErr

/-- `get <prop>(<inst>)` — evaluate a property applied to an instance.
    Parses "prop(inst1, inst2)" → `#eval prop inst1 inst2`. -/
def getVal (stateRef : IO.Ref KBState) (expr : String) : IO CmdResult := do
  let state ← stateRef.get
  let evalStr :=
    let byParen := expr.splitOn "("
    if byParen.length >= 2 then
      let func := (byParen.headD "").trim
      let args := ((byParen.tail |> String.intercalate "(").splitOn ")").headD "" |>.trim
      let argList := args.splitOn "," |>.map String.trim |> String.intercalate " "
      s!"{func} {argList}"
    else expr
  let (code, out) ← runLeanSnippet state s!"#eval {evalStr}"
  if code == 0 then println out.trim; return cmdOk
  else do println s!"Error evaluating '{expr}':\n{out.trim}"; return cmdErr

/-- Enhanced query with optional user-supplied tactic and proof term output.
    Syntax: query <prop> [using <tactic-script>]
    Default tactics: first | decide | aesop | simp
    On success: prints "Status: Yes" and the proof term via #print.
    On failure: prints "Status: No" or "Status: Unknown" with the Lean error. -/
def queryProp (stateRef : IO.Ref KBState) (prop : String) (tactics : Option String)
    : IO CmdResult := do
  let state  ← stateRef.get
  let tacs   := tactics.getD "first | decide | aesop | simp"
  let pname  := "_kb_proof"
  let thm    := s!"set_option maxHeartbeats 8000 in\ntheorem {pname} : {prop} := by {tacs}"
  let (code, out) ← runLeanSnippet state thm
  if code == 0 then
    -- Also fetch the proof term.
    let (_, printOut) ← runLeanSnippet state (thm ++ s!"\n#print {pname}")
    println s!"Status: Yes"
    println s!"Tactic: {tacs}"
    let proofLine := printOut.trim
    unless proofLine.isEmpty do println s!"Proof:\n{proofLine}"
    return cmdOk
  else if strContains out "false" || strContains out "False" ||
          strContains out "contradiction" then do
    println s!"Status: No"
    println s!"Tactic: {tacs}"
    println s!"Lean output:\n{out.trim}"
    return cmdErr
  else do
    println s!"Status: Unknown"
    println s!"Tactic: {tacs}"
    println s!"Lean output:\n{out.trim}"
    return cmdErr

--------------------------------------------------------------------------------
-- 12. Build / Commit / Status
--------------------------------------------------------------------------------

def existingDomains (root : System.FilePath) : IO (List String) := do
  let reserved := ["KnowledgeBase", "lakefile", "_kb_tmp"]
  let entries  ← System.FilePath.readDir root
  return dedupStrings <| entries.toList.filterMap fun e =>
    let fname := e.fileName
    if fname.endsWith ".lean" then
      let stem := fname.dropRight ".lean".length
      if reserved.contains stem then none else some stem
    else none

def syncRootModule (root : System.FilePath) (state : KBState) : IO Unit := do
  let allDomains ← existingDomains root
  for domain in allDomains do
    let ownDecls := state.decls.filter (fun d => d.domain == domain)
    let others   := allDomains.filter (· != domain)
    let mut needed : List String := []
    for other in others do
      let otherDecls := state.decls.filter (fun d => d.domain == other)
      if ownDecls.any fun od => otherDecls.any fun od2 => strContains od.lean od2.name then
        needed := needed ++ [other]
    if !needed.isEmpty then
      let path := domainPath root domain
      if ← path.pathExists then
        let contents ← IO.FS.readFile path
        let missing := needed.filter (fun imp => !strContains contents s!"import {imp}")
        unless missing.isEmpty do
          let header := missing.map (fun imp => s!"import {imp}") |> String.intercalate "\n"
          IO.FS.writeFile path (header ++ "\n\n" ++ contents)
  let mut root' := "-- KnowledgeBase root module (auto-generated)\n\n"
  for stem in allDomains do root' := root' ++ s!"import {stem}\n"
  IO.FS.writeFile (root / "KnowledgeBase.lean") root'

def checkKB (stateRef : IO.Ref KBState) : IO CmdResult := do
  let state ← stateRef.get
  syncRootModule state.root state
  println "Running `lake build`…"
  let code ← runProcPrint "lake" #["build"] (some state.root.toString)
  if code == 0 then println "✓  Consistent (build succeeded)."; return cmdOk
  else do println "✗  Build failed — see errors above."; return cmdErr

def commitChanges (stateRef : IO.Ref KBState) (msg : String) : IO CmdResult := do
  let state ← stateRef.get
  let dir := state.root.toString
  syncRootModule state.root state
  println "Verifying before commit…"
  let buildCode ← runProcPrint "lake" #["build"] (some dir)
  if buildCode != 0 then
    println "Commit aborted: build failed."; return cmdErr
  let _ ← runProcPrint "git" #["add", "."] (some dir)
  let code ← runProcPrint "git" #["commit", "-m", msg] (some dir)
  if code == 0 then println s!"Committed: \"{msg}\""; return cmdOk
  else do println "(Nothing to commit — working tree clean.)"; return cmdOk

def showStatus (stateRef : IO.Ref KBState) : IO CmdResult := do
  let state ← stateRef.get
  println "=== git status ==="
  let _ ← runProcPrint "git" #["status", "--short"] (some state.root.toString)
  println "\n=== knowledge base ==="
  println s!"  root  : {state.root}"
  println s!"  decls : {state.decls.size} loaded"
  let domains := dedupStrings (state.decls.map (·.domain)).toList
  for dom in domains do
    let n := (state.decls.filter (fun d => d.domain == dom)).size
    println s!"  {dom}: {n}"
  return cmdOk

--------------------------------------------------------------------------------
-- 13. Init
--------------------------------------------------------------------------------

def initKB (dir : String) : IO CmdResult := do
  let root : System.FilePath := System.FilePath.mk dir
  if (← runProcPrint "mkdir" #["-p", dir]) != 0 then
    println s!"Error: could not create '{dir}'."; return cmdErr
  IO.FS.writeFile (root / "lakefile.lean")
    "import Lake\nopen Lake DSL\n\npackage KnowledgeBase where\n\nlean_lib KnowledgeBase where\n"
  IO.FS.writeFile (root / "KnowledgeBase.lean")
    "-- KnowledgeBase root module\n-- Run `kb check` to populate this file.\n"
  let _ ← runProcPrint "git" #["init"]                                (some dir)
  let _ ← runProcPrint "git" #["add", "."]                            (some dir)
  let _ ← runProcPrint "git" #["commit", "-m", "Initial kb scaffold"] (some dir)
  println s!"Initialised knowledge base in '{dir}'."
  return cmdOk

--------------------------------------------------------------------------------
-- 14. Domain info
--------------------------------------------------------------------------------

def listDomains (stateRef : IO.Ref KBState) : IO CmdResult := do
  let state   ← stateRef.get
  let domains := dedupStrings (state.decls.map (·.domain)).toList
  if domains.isEmpty then println "(no domains loaded — use 'add … in <Domain>' to create one)"
  else do
    println "Domains:"
    for dom in domains do
      let n := (state.decls.filter (fun d => d.domain == dom)).size
      println s!"  {dom}  ({n} declaration(s))"
  return cmdOk

def domainsStats (stateRef : IO.Ref KBState) : IO CmdResult := do
  let state   ← stateRef.get
  let domains := dedupStrings (state.decls.map (·.domain)).toList
  println "=== Domain statistics ==="
  if domains.isEmpty then println "  (none)"
  else for dom in domains do
    let ds := state.decls.filter (fun d => d.domain == dom)
    println s!"  {dom}: {ds.size}"
    for k in ["type","instance","prop","rel","fact","rule","def"] do
      let n := (ds.filter (fun d => d.kind == k)).size
      if n > 0 then println s!"    {k}: {n}"
  println "\n=== By kind ==="
  for k in ["type","instance","prop","rel","fact","rule","def"] do
    let n := (state.decls.filter (fun d => d.kind == k)).size
    if n > 0 then println s!"  {k}: {n}"
  println s!"\n  Total: {state.decls.size}"
  return cmdOk

--------------------------------------------------------------------------------
-- 15. Help
--------------------------------------------------------------------------------

def helpText : String :=
  "Lean4 Knowledge Base CLI (kb)\n" ++
  "All `add` commands require `in <Domain>` — each domain is a <Domain>.lean file.\n\n" ++
  "ADD\n" ++
  "  add type <n> [extends <Parent>] in <D>\n" ++
  "  add instance <n> : <Type> in <D>\n" ++
  "  add prop <n> : <DomType> -> <Range> in <D>\n" ++
  "  add rel <n> : <D1> -> <D2> -> Prop in <D>\n" ++
  "  add fact <fact> in <D>          prop(inst) = val  or  rel(a, b)\n" ++
  "  add rule <n> : <statement> in <D>\n" ++
  "  add def <n> <params> : <Type> := <body> in <D>\n\n" ++
  "READ\n" ++
  "  list <types|instances|props|rels|facts|rules|defs|all> [of <Domain>]\n" ++
  "  show def <n>           Lean source of a declaration.\n" ++
  "  get <prop>(<inst>)     Retrieve a property value (runs #eval).\n" ++
  "  search <text>          Full-text search across names, Lean source, domains.\n" ++
  "  deps <n>               Forward and reverse dependency graph for a declaration.\n" ++
  "  eval <expr>            Run #eval on a Lean expression.\n\n" ++
  "QUERY\n" ++
  "  query <prop>\n" ++
  "  query <prop> using <tactic-script>\n" ++
  "      Attempt to prove a proposition.  Default tactics: first | decide | aesop | simp\n" ++
  "      Returns Status: Yes / No / Unknown and the proof term on success.\n\n" ++
  "BUILD & GIT\n" ++
  "  check                  Sync imports + lake build to verify whole KB.\n" ++
  "  commit <message>       Verify then git-commit all changes.\n" ++
  "  status                 git status + domain summary.\n" ++
  "  rm <n>                 Remove a declaration (warns about dependents).\n\n" ++
  "DOMAINS\n" ++
  "  domains                List domains with declaration counts.\n" ++
  "  domains stats          Breakdown per domain and kind.\n\n" ++
  "  init <dir>             Create a new KB repository.\n" ++
  "  help | exit\n\n" ++
  "EXIT CODES  0 = success   1 = error\n\n" ++
  "PERSISTENCE  State is stored as --@kb metadata in each domain file and\n" ++
  "             reloaded automatically on startup.\n\n" ++
  "EXAMPLE\n" ++
  "  add type Entity in People\n" ++
  "  add type Person extends Entity in People\n" ++
  "  add instance socrates : Person in People\n" ++
  "  add prop birthYear : Person -> Nat in People\n" ++
  "  add fact birthYear(socrates) = 470 in People\n" ++
  "  get birthYear(socrates)\n" ++
  "  deps birthYear\n" ++
  "  search socrates\n" ++
  "  query birthYear socrates = 470 using decide\n" ++
  "  check\n" ++
  "  commit \"Added Socrates\"\n"

--------------------------------------------------------------------------------
-- 16. Main REPL loop
--------------------------------------------------------------------------------

partial def processCommand (stateRef : IO.Ref KBState) (line : String) : IO CmdResult := do
  match parseCommand line with
  | none =>
    println s!"Error: unknown command '{line}'."
    println "  Hint: add commands require `in <Domain>`.  Run 'help' for usage."
    return cmdErr
  | some cmd =>
    match cmd with
    | .Init dir              => initKB dir
    | .AddType n p dom       => addType     stateRef n p dom
    | .AddInstance n t dom   => addInstance stateRef n t dom
    | .AddProp n dt rng dom  => addProp     stateRef n dt rng dom
    | .AddRel n d1 d2 dom    => addRel      stateRef n d1 d2 dom
    | .AddFact f dom         => addFact     stateRef f dom
    | .AddRule n s dom       => addRule     stateRef n s dom
    | .AddDef n p t b dom    => addDef      stateRef n p t b dom
    | .Remove n              => removeDecl  stateRef n
    | .List w f              => listDecls   stateRef w f
    | .ShowDef n             => showDef     stateRef n
    | .Get e                 => getVal      stateRef e
    | .Search q              => searchDecls stateRef q
    | .Deps n                => showDeps    stateRef n
    | .Eval e                => evalExpr    stateRef e
    | .Query p t             => queryProp   stateRef p t
    | .Check                 => checkKB     stateRef
    | .Commit m              => commitChanges stateRef m
    | .Status                => showStatus  stateRef
    | .Domains               => listDomains stateRef
    | .DomainsStats          => domainsStats stateRef
    | .Help                  => println helpText; return cmdOk
    | .Exit                  => return cmdQuit

partial def main : IO Unit := do
  let root  : System.FilePath := System.FilePath.mk "."
  let decls ← loadState root
  let stateRef ← IO.mkRef (KBState.mk root decls)
  if decls.size > 0 then
    println s!"Loaded {decls.size} declaration(s) from KB."
  println "Lean4 Knowledge Base CLI — type 'help' for commands, 'exit' to quit."
  let rec loop : IO Unit := do
    IO.print "kb> "
    let line ← (← IO.getStdin).getLine
    let trimmed := line.trim
    unless trimmed.isEmpty do
      let result ← processCommand stateRef trimmed
      unless result.continue do return
    loop
  loop