codegen.lean

import Lean.Meta
import LeanAide.Responses
import LeanAide.PaperCodes
import LeanAide.Codegen
import LeanAideCore.Translate
import Mathlib

open Lean LeanAide

set_option maxHeartbeats 10000000
set_option maxRecDepth 1000
set_option compiler.extract_closed false

def usage : String :=
  "Usage: lake env lean --run codegen.lean <structured-json-file>"

def stripJsonSuffix (name : String) : String :=
  if name.endsWith ".json" then
    (name.dropEnd ".json".length).toString
  else
    name

def outputPathFor (inputPath : System.FilePath) : System.FilePath :=
  let inputName := inputPath.fileName.getD "generated.json"
  let stem := stripJsonSuffix inputName
  System.mkFilePath ["CodeGen", stem ++ ".lean"]

def taskJson (js : Json) : Json :=
  match js.getObjVal? "document_json" with
  | Except.ok _ => js
  | Except.error _ => Json.mkObj [("document_json", js)]

def ensureMathlibImport (topCode : String) : String :=
  if topCode.contains "import Mathlib" then
    topCode
  else
    "import Mathlib\n" ++ topCode

def generatedFileContents (inputPath : System.FilePath) (topCode documentCode : String) : String :=
  let topCode := ensureMathlibImport topCode
  s!"{topCode.trimAsciiEnd}\n\n-- Generated by codegen.lean from {inputPath}\n\n{documentCode.trimAsciiEnd}\n"

def countLabel (n : Nat) (singular plural : String) : String :=
  if n == 1 then s!"1 {singular}" else s!"{n} {plural}"

def indentLines (indent : String) (text : String) : String :=
  String.intercalate "\n" <| text.splitOn "\n" |>.map (indent ++ ·)

def jsonFieldAsString (js : Json) (field : String) : String :=
  match js.getObjValAs? String field with
  | Except.ok value => value
  | Except.error _ =>
      match js.getObjVal? field with
      | Except.ok value => value.compress
      | Except.error _ => "<missing>"

def flattenJsonArrayItems (items : Array Json) : Array Json :=
  items.foldl
    (fun acc item =>
      match item with
      | Json.arr nested => acc ++ nested
      | _ => acc.push item)
    #[]

def jsonObjArrayField (js : Json) (field : String) : Except String (Array Json) := do
  let items ← js.getObjValAs? (Array Json) field
  pure <| flattenJsonArrayItems items

def formatStringList (title empty : String) (items : Array String) : String :=
  if items.isEmpty then
    empty
  else
    let numbered := items.toList.zipIdx.map fun (item, idx) =>
      s!"{idx + 1}. {indentLines "   " item |>.drop 3}"
    s!"{title} ({items.size}):\n{String.intercalate "\n" numbered}"

def formatSorryItem (idx : Nat) (js : Json) : String :=
  let declaration := jsonFieldAsString js "declaration_name"
  let sorryType := jsonFieldAsString js "sorry_type"
  let header := s!"{idx + 1}. {declaration}"
  s!"{header}\n{indentLines "   " sorryType}"

def formatSorries (title empty : String) (items : Array Json) : String :=
  if items.isEmpty then
    empty
  else
    let rendered := items.toList.zipIdx.map fun (item, idx) => formatSorryItem idx item
    s!"{title} ({countLabel items.size "goal" "goals"}):\n{String.intercalate "\n\n" rendered}"

def formatElaborationResult (result : Json) : String :=
  let declarations :=
    result.getObjValAs? (Array String) "declarations" |>.toOption.getD #[]
  let logs :=
    result.getObjValAs? (Array String) "logs" |>.toOption.getD #[]
  let sorries :=
    jsonObjArrayField result "sorries" |>.toOption.getD #[]
  let sorriesAfterPurge :=
    jsonObjArrayField result "sorries_after_purge" |>.toOption.getD #[]
  let declSection :=
    if declarations.isEmpty then
      "Declarations: none"
    else
      s!"Declarations ({declarations.size}): {String.intercalate ", " declarations.toList}"
  String.intercalate "\n\n" [
    declSection,
    formatStringList "Elaboration logs" "Elaboration logs: none" logs,
    formatSorries "Sorries" "Sorries: none" sorries,
    formatSorries "Sorries after purge" "Sorries after purge: none" sorriesAfterPurge
  ]

def runCodegen (inputPath : System.FilePath) : IO UInt32 := do
  initSearchPath (← findSysroot)
  let raw ← IO.FS.readFile inputPath
  let js ←
    match Json.parse raw with
    | Except.ok js => pure js
    | Except.error err => do
      IO.eprintln s!"Failed to parse JSON from {inputPath}: {err}"
      return 1
  let env ←
    importModules (loadExts := true) #[
      {module := `Mathlib},
      {module := `LeanAideCore.Translate},
      {module := `LeanAide.PaperCodes},
      {module := `LeanAide.Responses},
      {module := `LeanAide.Codegen}] {}
  let ctx : Core.Context :=
    { fileName := inputPath.toString
      fileMap := {source := raw, positions := #[]}
      maxHeartbeats := 0
      maxRecDepth := 1000000 }
  let translator : Translator := {}
  let core := leanFromStructuredJsonTask (taskJson js) translator |>.runToCore
  let result ← core.run' ctx {env := env} |>.runToIO'
  match result.getObjValAs? String "result" with
  | Except.ok "success" =>
      let documentCode ←
        match result.getObjValAs? String "document_code" with
        | Except.ok code => pure code
        | Except.error err => do
          IO.eprintln s!"Code generation succeeded but no document_code was returned: {err}"
          return 1
      let topCode := result.getObjValAs? String "top_code" |>.toOption.getD "import Mathlib\n"
      let outFile := outputPathFor inputPath
      IO.FS.createDirAll (System.mkFilePath ["CodeGen"])
      IO.FS.writeFile outFile (generatedFileContents inputPath topCode documentCode)
      IO.println s!"Generated {outFile}"
      match result.getObjVal? "declarations" with
      | Except.ok decls => IO.eprintln s!"Declarations: {decls.compress}"
      | Except.error _ => pure ()
      let core := elaborateTask result translator |>.runToCore
      let result' ← core.run' ctx {env := env} |>.runToIO'
      match result'.getObjValAs? String "result" with
      | Except.ok "success" =>
        IO.eprintln "Elaboration succeeded"
        IO.eprintln (formatElaborationResult result')
        let handle ← IO.FS.Handle.mk outFile IO.FS.Mode.append
        handle.putStrLn "\n\n/-!\n## Elaboration results \n\n"
        handle.putStrLn (formatElaborationResult result')
        handle.putStrLn "\n\n-/\n"
        handle.flush
        return 0
      | Except.ok other => do
          IO.eprintln s!"Elaboration returned non-success result: {other}"
          IO.eprintln (formatElaborationResult result')
          let handle ← IO.FS.Handle.mk outFile IO.FS.Mode.append
          handle.putStrLn "\n\n/-!\n## Elaboration results \n\n"
          handle.putStrLn (formatElaborationResult result')
          handle.putStrLn "\n\n-/\n"
          handle.flush
          return 1
      | Except.error err => do
          IO.eprintln s!"Elaboration returned malformed response: {err}"
          IO.eprintln result'.pretty
          return 1
  | Except.ok other => do
      IO.eprintln s!"Code generation returned non-success result: {other}"
      IO.eprintln result.pretty
      return 1
  | Except.error err => do
      IO.eprintln s!"Code generation returned malformed response: {err}"
      IO.eprintln result.pretty
      return 1

def main (args : List String) : IO UInt32 := do
  match args with
  | [input] => runCodegen input
  | _ => do
      IO.eprintln usage
      return 1