/-
This is the Lean 4 game server. It offers a way to interact with Lean 4 which
is completely distinct from running the command-line lean or the language server protocol.

It is based on lean-gym by Daniel Selsam.
-/
import Lean.Data.Json.Basic

import GameServer.Utils
import GameServer.EnvExtensions
import GameServer.Communication

open Lean Meta Elab Tactic Std

/-- Convert format to string without line breaks -/
def Std.Format.oneline : Format → String
| .nil => ""
| .line => ""
| .text s => s
| .nest _ f => f.oneline
| .append f g => f.oneline ++ g.oneline
| .group f _ => f.oneline
| .tag _ f => f.oneline

/-- Convert JSON to string without line breaks -/
instance instToStringJsonOneLine : ToString Json := ⟨fun o => o.render.oneline⟩
attribute [-instance] Lean.Json.instToStringJson

/-! ## GameGoal -/
structure GameGoal where
  objects : List LocalDecl
  assumptions : List LocalDecl
  goal : String
  mvarid : MVarId

def Lean.MVarId.toGameGoal (goal : MVarId) : MetaM GameGoal := do
match (← getMCtx).findDecl? goal with
| none          => throwError "unknown goal"
| some mvarDecl => do
  -- toGameGoalAux below will sort local declarations from the context of goal into data and assumptions,
  -- discarding auxilliary declarations
  let rec toGameGoalAux : List (Option LocalDecl) → MetaM (List LocalDecl × List LocalDecl)
  | (some decl)::t => withLCtx mvarDecl.lctx mvarDecl.localInstances do
    let (o, a) ← toGameGoalAux t
    if decl.isAuxDecl then
      return (o, a)
    if (← inferType decl.type).isProp then
      return (o, decl::a)
    else
      return (decl::o, a)
  | none:: t => toGameGoalAux t
  | [] => return ([], [])
  withLCtx mvarDecl.lctx mvarDecl.localInstances do
    let (objects, assumptions) ← toGameGoalAux mvarDecl.lctx.decls.toList
    return {objects := objects, assumptions := assumptions, goal := toString (← Meta.ppExpr mvarDecl.type),
            mvarid := goal }

def GameGoal.toJson (gameGoal : GameGoal) : MetaM Json :=
gameGoal.mvarid.withContext do
  return Json.mkObj [("objects", Lean.ToJson.toJson (← gameGoal.objects.mapM Lean.LocalDecl.toJson)),
                     ("assumptions", Lean.ToJson.toJson (← gameGoal.assumptions.mapM Lean.LocalDecl.toJson)),
                     ("goal", gameGoal.goal)]


/-! ## Action -/

inductive Action where
| info : Action
| loadLevel : Nat → Action
| runTactic : String → Action
| undo : Action
| restartGame : Action
| restartLevel : Action
| quit : Action
| next : Action
| prev : Action
| invalid : String → Action  -- Used for broken Json parsing
deriving ToJson, FromJson, Repr

def Action.parse (s : String) : Action :=
  let e : Except String Action :=
    try
      return ← fromJson? (← Json.parse s)
    catch _ => return Action.invalid s
  match e with
  | Except.ok a => a
  | _ => Action.info

def Action.get : IO Action :=
  return Action.parse (← (← IO.getStdin).getLine)


/-! ## LevelInfo -/

structure LevelInfo extends GameLevel where
  goals    : List GameGoal := []
  errors   : Array String := #[]
  message  : String := ""

def LevelInfo.toJson (info : LevelInfo) : MetaM Json :=
  return Json.mkObj [("index", Lean.ToJson.toJson info.index),
                     ("title", Lean.ToJson.toJson info.title),
                     ("tactics", Lean.ToJson.toJson info.tactics),
                     ("lemmas", Lean.ToJson.toJson info.lemmas),
                     ("errors", Lean.ToJson.toJson info.errors),
                     ("goals", Lean.ToJson.toJson (← info.goals.mapM (·.toJson))),
                     ("message", info.message)]

namespace Server

/-! ## LevelM and Response -/

structure SavedState where
  tacticState : Tactic.SavedState
  message : String := ""

abbrev LevelM := ReaderT GameLevel StateRefT (Array SavedState) TermElabM

/-- Returns the most recent SavedState. -/
def getState : LevelM SavedState := do
  let some state := (← get).back? | throwError "Couldn't find tactic state"
  return state


structure Response : Type where
  goals    : List GameGoal := []
  errors   : Array String := #[]
  message  : String := ""

def Response.toJson (resp : Response) : MetaM Json :=
  return Json.mkObj [("errors", Lean.ToJson.toJson resp.errors),
                     ("goals", Lean.ToJson.toJson (← resp.goals.mapM (·.toJson))),
                     ("message", resp.message)]

/-! ## Main running functions -/

/-- Dummy `Core.Context` value to be fed to `Lean.Core.CoreM.toIO` -/
def coreCtx : Core.Context := {
  currNamespace := Name.anonymous,
  openDecls := [],
  fileName := "<Game>",
  fileMap := { source := "", positions := #[0], lines := #[1] } }


structure RunTacticParams where
  tactic : String
deriving FromJson

structure LoadLevelParams where
  number : Nat
deriving FromJson

partial def runLevel (requestId : JsonRpc.RequestID) (env : Environment) (GameName : Name) (levels : HashMap Nat GameLevel) (idx : Nat) : IO Unit := do
  let levelName : Name := s!"Level{toString idx}"
  let termElabM : TermElabM Unit := do
    let some lvl := levels.find? idx | throwError s!"Cannot find level {idx}"
    let mvar ← mkFreshExprMVar (some lvl.goal) (kind := MetavarKind.synthetic)
    let (_, mvar) ← mvar.mvarId!.introNP lvl.intro_nb
    mvar.withContext do
      let state := #[{tacticState := { term := ← Term.saveState, tactic := { goals := [mvar] }}}]
      let levelM : LevelM Unit := do
          let resp := {← mkResponse with message := lvl.introduction}
          let levelInfo : LevelInfo :=
              { index := lvl.index,
                title := lvl.title,
                tactics := lvl.tactics,
                lemmas := lvl.lemmas,
                errors := resp.errors,
                goals := resp.goals,
                message := resp.message }
          let hOut ← IO.getStdout
          hOut.writeGspResponse ⟨requestId, ← levelInfo.toJson⟩
          mainLoop
      levelM.run lvl |>.run' state
  let metaM : MetaM Unit := termElabM.run' (ctx := {})
  try
    discard <| metaM.run'.toIO coreCtx { env := env }
  catch
  | .userError s => output s!"Could not run level {idx}: {s}"
  | _ => output s!"Could not run level {idx}"

where

  mkResponse (errors : Array String := #[]) : LevelM Response := do
    let savedState ← getState
    let goals := savedState.tacticState.tactic.goals
    return { goals := (← liftM $ goals.mapM Lean.MVarId.toGameGoal), message := savedState.message, errors := errors }

  /-- Try to parse the given String as a single line tactic invocation. Update the LevelM state only
  if the tactic succeeds. Always return a Response object which contains information about the current
  state with a message and errors if any. -/
  runTactic (tacticString : String) : LevelM Response := do
    let lvl ← read
    let tacticNames := (lvl.tactics.map (·.name.toString)).toList
    let tacName := (tacticString.trim.splitOn " ").headD ""
    if not (tacticNames.contains tacName) then mkResponse #["Unrecognized tactic"] else
      let savedState ← getState
      match Parser.runParserCategory (← getEnv) `tactic tacticString "<stdin>" with
      | Except.error err => mkResponse #[err]
      | Except.ok stx    => do
        savedState.tacticState.term.restore
        let mvarId : MVarId := savedState.tacticState.tactic.goals.head!
        try
          let unsolvedGoals ← Tactic.run mvarId do set savedState.tacticState.tactic
                                                   evalTactic stx
          if (← getThe Core.State).messages.hasErrors then
            let messages := (← getThe Core.State).messages.getErrorMessages.toList.toArray
            mkResponse (← (messages.map Message.data).mapM fun md => md.toString)
          else
            let savedState : Tactic.SavedState := { term := (← Term.saveState), tactic := { goals := unsolvedGoals}}
            let mut message := ""
            match unsolvedGoals with
            | goal::_ => do
              let normalized_tgt ← normalizedRevertExpr goal
              for msg in lvl.messages do
                if ← isDefEq normalized_tgt msg.normalized_goal then
                  message := msg.message
                  break
            | [] => pure ()
            if unsolvedGoals matches [] then
              message := (← read).conclusion
            modify fun s => s.push {tacticState := savedState, message := message}
            mkResponse
        catch ex => mkResponse #[← ex.toMessageData.toString]

  mainLoop : LevelM Unit := do
    let hIn ← IO.getStdin
    let hOut ← IO.getStdout
    let hLog ← IO.getStderr
    let m ← hIn.readGspMessage
    match m with
    | .request id "runTactic" params => do
      match fromJson? (toJson (params)) with
      | Except.ok (v : RunTacticParams) =>
        let resp ← runTactic v.tactic
        hOut.writeGspResponse ⟨id, ← resp.toJson⟩
      | Except.error inner =>
        hLog.putStr s!"Invalid params: {inner}"
        hLog.flush
    | .request id "loadLevel" (some params) =>
      match fromJson? (toJson (params)) with
      | Except.ok (v : LoadLevelParams) =>
        runLevel id env GameName levels v.number
      | Except.error inner =>
        hLog.putStr s!"Invalid params: {inner}"
        hLog.flush
    | .request id "undo" _ => do
       modify fun s => s.pop
       hOut.writeGspResponse ⟨id, ← (← mkResponse).toJson⟩
    | _ =>
      hLog.putStr s!"Invalid action: {toJson m}"
      hLog.flush
    mainLoop

open System Lean Std in
partial def runGame (GameName : Name) : IO Unit := do
  let env ← importModules [{ module := `Init : Import }, { module := GameName : Import }] {} 0
  let termElabM : TermElabM Unit := do
    let levels := levelsExt.getState env
    let game := {← gameExt.get with nb_levels := levels.size }
    mainLoop env game levels
  let metaM : MetaM Unit := termElabM.run' (ctx := {})
  discard <| metaM.run'.toIO coreCtx { env := env }
where
  mainLoop (env : Environment) (game : Game) (levels : HashMap Nat GameLevel): IO Unit := do
    let hIn ← IO.getStdin
    let hOut ← IO.getStdout
    let hLog ← IO.getStderr
    let m ← hIn.readGspMessage
    match m with
    | .request id "info" params =>
      hOut.writeGspResponse ⟨id, game⟩
    | .request id "loadLevel" (some params) =>
      match fromJson? (toJson (params)) with
      | Except.ok (v : LoadLevelParams) =>
        runLevel id env GameName levels v.number
      | Except.error inner =>
        hLog.putStr s!"Invalid params: {inner}"
        hLog.flush
    | .request id m _ =>
      hOut.writeGspResponseError
        ⟨id, JsonRpc.ErrorCode.methodNotFound, s!"Method not found: {m}", none⟩
    | _ =>
      hLog.putStr s!"Invalid action: {toJson m}"
      hLog.flush

    mainLoop env game levels

end Server