lean4game/server/GameServer/RpcHandlers.lean

import GameServer.EnvExtensions
import GameServer.InteractiveGoal

open Lean
open Server
open Widget
open RequestM
open Meta

/-! ## GameGoal -/

namespace GameServer

def levelIdFromFileName? (initParams : Lsp.InitializeParams) (fileName : String) : Option LevelId := Id.run do
  let fileParts := fileName.splitOn "/"
  if fileParts.length == 3 then
    if let (some level, some game) := (fileParts[2]!.toNat?, initParams.rootUri?) then
      return some {game, world := fileParts[1]!, level := level}
  return none

def getLevelByFileName? [Monad m] [MonadEnv m] (initParams : Lsp.InitializeParams) (fileName : String) : m (Option GameLevel) := do
  let some levelId := levelIdFromFileName? initParams fileName
    | return none
  return ← getLevel? levelId

structure FVarBijection :=
  (forward : HashMap FVarId FVarId)
  (backward : HashMap FVarId FVarId)

instance : EmptyCollection FVarBijection := ⟨{},{}⟩

def FVarBijection.insert (bij : FVarBijection) (a b : FVarId) : FVarBijection :=
  ⟨bij.forward.insert a b, bij.backward.insert b a⟩

def FVarBijection.insert? (bij : FVarBijection) (a b : FVarId) : Option FVarBijection :=
  let a' := bij.forward.find? a
  let b' := bij.forward.find? b
  if (a' == none || a' == some b) && (b' == none || b' == some a)
  then some $ bij.insert a b
  else none

/-- Checks if `pattern` and `e` are equal up to FVar identities. -/
partial def matchExpr (pattern : Expr) (e : Expr) (bij : FVarBijection := {}) : Option FVarBijection :=
  match pattern, e with
  | .bvar i1, .bvar i2 => if i1 == i2 then bij else none
  | .fvar i1, .fvar i2 => bij.insert? i1 i2
  | .mvar _, .mvar _ => bij
  | .sort _u1, .sort _u2 => bij -- TODO?
  | .const n1 _ls1, .const n2 _ls2 =>
    if n1 == n2 then bij else none -- && (← (ls1.zip ls2).allM fun (l1, l2) => Meta.isLevelDefEq l1 l2)
  | .app f1 a1, .app f2 a2 =>
    some bij
      |> (Option.bind · (fun bij => matchExpr f1 f2 bij))
      |> (Option.bind · (fun bij => matchExpr a1 a2 bij))
  | .lam _ t1 b1 _, .lam _ t2 b2 _ =>
    some bij
      |> (Option.bind · (fun bij => matchExpr t1 t2 bij))
      |> (Option.bind · (fun bij => matchExpr b1 b2 bij))
  | .forallE _ t1 b1 _, .forallE _ t2 b2 _ =>
    some bij
      |> (Option.bind · (fun bij => matchExpr t1 t2 bij))
      |> (Option.bind · (fun bij => matchExpr b1 b2 bij))
  | .letE _ t1 v1 b1 _, .letE _ t2 v2 b2 _ =>
    some bij
      |> (Option.bind · (fun bij => matchExpr t1 t2 bij))
      |> (Option.bind · (fun bij => matchExpr v1 v2 bij))
      |> (Option.bind · (fun bij => matchExpr b1 b2 bij))
  | .lit l1, .lit l2 =>
    if l1 == l2 then bij else none
  | .proj i1 n1 e1, .proj i2 n2 e2 =>
    if i1 == i2 && n1 == n2 then matchExpr e1 e2 bij else none
  -- ignore mdata:
  | .mdata _ pattern', _ =>
    matchExpr pattern' e bij
  | _, .mdata _ e' =>
    matchExpr pattern e' bij
  | _, _ => none

/-- Check if each fvar in `patterns` has a matching fvar in `fvars` -/
def matchDecls (patterns : Array Expr) (fvars : Array Expr) (strict := true) (initBij : FVarBijection := {}) : MetaM (Option FVarBijection) := do
  -- We iterate through the array backwards hoping that this will find us faster results
  -- TODO: implement backtracking
  let mut bij := initBij
  for i in [:patterns.size] do
    let pattern := patterns[patterns.size - i - 1]!
    if bij.forward.contains pattern.fvarId! then
      continue
    for j in [:fvars.size] do
      let fvar := fvars[fvars.size - j - 1]!
      if bij.backward.contains fvar.fvarId! then
        continue

      if let some bij' := matchExpr
          (← instantiateMVars $ ← inferType pattern)
          (← instantiateMVars $ ← inferType fvar) bij then
        -- usedFvars := usedFvars.set! (fvars.size - j - 1) true
        bij := bij'.insert pattern.fvarId! fvar.fvarId!
        break
    if ! bij.forward.contains pattern.fvarId! then return none

  if !strict || fvars.all (fun fvar => bij.backward.contains fvar.fvarId!)
  then return some bij
  else return none

unsafe def evalHintMessageUnsafe : Expr → MetaM (Array Expr → MessageData) :=
  evalExpr (Array Expr → MessageData)
    (.forallE default (mkApp (mkConst ``Array [levelZero]) (mkConst ``Expr))
      (mkConst ``MessageData) .default)

@[implemented_by evalHintMessageUnsafe]
def evalHintMessage : Expr → MetaM (Array Expr → MessageData) := fun _ => pure (fun _ => "")

open Meta in
/-- Find all hints whose trigger matches the current goal -/
def findHints (goal : MVarId) (m : DocumentMeta) (initParams : Lsp.InitializeParams) : MetaM (Array GameHint) := do
  goal.withContext do
    let some level ← getLevelByFileName? initParams m.mkInputContext.fileName
      | throwError "Level not found: {m.mkInputContext.fileName}"
    let hints ← level.hints.filterMapM fun hint => do
      openAbstractCtxResult hint.goal fun hintFVars hintGoal => do
        if let some fvarBij := matchExpr (← instantiateMVars $ hintGoal) (← instantiateMVars $ ← inferType $ mkMVar goal)
        then
          let lctx := (← goal.getDecl).lctx
          if let some bij ← matchDecls hintFVars lctx.getFVars (strict := hint.strict) (initBij := fvarBij)
          then
            let userFVars := hintFVars.map fun v => bij.forward.findD v.fvarId! v.fvarId!
            let text := (← evalHintMessage hint.text) (userFVars.map Expr.fvar)
            let ctx := {env := ← getEnv, mctx := ← getMCtx, lctx := lctx, opts := {}}
            let text ← (MessageData.withContext ctx text).toString
            return some { text := text, hidden := hint.hidden }
          else return none
        else
          return none
    return hints

/-- Get the line that ends in `pos`. Note that `pos` is expected to be the
position of a `\n` but this is not enforced. -/
def _root_.Lean.FileMap.getLineBefore (fmap : FileMap) (pos : String.Pos) : String := Id.run do
  match fmap.positions.findIdx? (· == pos) with
  | none =>
    panic s!"Position {pos} is not a newline character in " ++
      s!"the following string: '{fmap.source}'!"
  | some 0 =>
    -- the first entry of `positions` is always `0`
    return ""
  | some (i + 1) =>
    let line : Substring := ⟨fmap.source, fmap.positions.get! i, pos⟩
    return line.toString

/-- Returns the `List` without the last element. -/
def _root_.List.dropBack {α : Type _} : List α → List α
| [] => []
| _ :: [] => []
| x :: xs => x :: xs.dropBack

/-- Trim empty lines from the file and add a single newline. -/
def _root_.Lean.FileMap.trim (fmap : FileMap) : FileMap :=
  let str := match fmap.source.trim with
  | "" => ""
  | s => s ++ "\n"
  FileMap.ofString str

/-- Returns the `Array` without the last element. -/
def _root_.Array.dropBack {α : Type _} (a : Array α) : Array α := ⟨a.data.dropBack⟩

/-- Add custom diagnostics about whether the level is completed. -/
def addCompletionDiagnostics (diag : Array InteractiveDiagnostic) (goals : Array InteractiveGoalWithHints)
    (pos : Lsp.Position) (prevGoalCount : Nat) : RequestM <| Array InteractiveDiagnostic := do
  let mut out : Array InteractiveDiagnostic := diag

  if goals.size == 0 then
    if diag.any (·.severity? == some .error) then
      pure ()
    else if diag.any (·.severity? == some .warning) then
      out := out.push {
        message := .text "level completed with warnings. 🎭"
        range := {
          start := pos
          «end» := pos
          }
        severity? := Lsp.DiagnosticSeverity.information }
    else
      out := out.push {
        message := .text "level completed! 🎉"
        range := {
          start := pos
          «end» := pos
          }
        severity? := Lsp.DiagnosticSeverity.information }
  else if goals.size < prevGoalCount then
    out := out.push {
      message := .text "intermediate goal solved! 🎉"
      range := {
        start := pos
        «end» := pos
        }
      severity? := Lsp.DiagnosticSeverity.information
    }


  return out
          -- diagsAtPos := if goalsAtPos.size < intermediateGoalCount then
          --   diagsAtPos.push {
          --     message := .text "intermediate goal solved 🎉"
          --     range := {
          --       start := lspPosAt
          --       «end» := lspPosAt
          --       }
          --     severity? := Lsp.DiagnosticSeverity.information
          --   }
          --   else diagsAtPos

/-- Request that returns the goals at the end of each line of the tactic proof
plus the diagnostics (i.e. warnings/errors) for the proof.
 -/
def getProofState (_ : Lsp.PlainGoalParams) : RequestM (RequestTask (Option ProofState)) := do
  let doc ← readDoc
  let rc ← readThe RequestContext
  let text := doc.meta.text.trim

  -- BUG: trimming here is a problem, since the snap might already be evaluated before
  -- the trimming and then the positions don't match anymore :((

  withWaitFindSnap
    doc
    -- TODO (Alex): I couldn't find a good condition to find the correct snap. So we are looking
    -- for the first snap with goals here.
    -- NOTE (Jon): The entire proof is in one snap, so hoped that Position `0` is good enough.
    (fun snap => ¬ (snap.infoTree.goalsAt? doc.meta.text 0).isEmpty)
    (notFoundX := return none)
    fun snap => do
      -- `snap` is the one snapshot containing the entire proof.
      let mut steps : Array <| InteractiveGoalsWithHints := #[]

      -- Question: Is there a difference between the diags of this snap and the last snap?
      -- Should we get the diags from there?
      -- Answer: The last snap only copied the diags from the end of this snap
      let mut diag : Array InteractiveDiagnostic := snap.interactiveDiags.toArray

      let mut intermediateGoalCount := 0

      -- Drop the last position as we ensured that there is always a newline at the end
      for pos in text.positions.dropBack do
        -- iterate over all newlines in the proof and get the goals and hints at each position
        -- TODO: we drop the last position because we always have a newline. Would be better
        -- to trim newlines instead before submitting!
        let source := text.getLineBefore pos

        if let goalsAtResult@(_ :: _) := snap.infoTree.goalsAt? doc.meta.text pos then
          pure ()
          let goalAtPos : List <| List InteractiveGoalWithHints ← goalsAtResult.mapM
            fun { ctxInfo := ci, tacticInfo := tacticInfo, useAfter := useAfter, .. } => do
              -- TODO: What does this function body do?
              -- let ciAfter := { ci with mctx := ti.mctxAfter }
              let ci := if useAfter then
                  { ci with mctx := tacticInfo.mctxAfter }
                else
                  { ci with mctx := tacticInfo.mctxBefore }
              -- compute the interactive goals
              let goalMvars : List MVarId ← ci.runMetaM {} do
                return if useAfter then tacticInfo.goalsAfter else tacticInfo.goalsBefore

              let interactiveGoals : List InteractiveGoalWithHints ← ci.runMetaM {} do
                goalMvars.mapM fun goal => do
                  let hints ← findHints goal doc.meta rc.initParams
                  let interactiveGoal ← goalToInteractive goal
                  return ⟨interactiveGoal, hints⟩
              -- TODO: This code is way old, can it be deleted?
              -- compute the goal diff
              -- let goals ← ciAfter.runMetaM {} (do
              --     try
              --       Widget.diffInteractiveGoals useAfter ti goals
              --     catch _ =>
              --       -- fail silently, since this is just a bonus feature
              --       return goals
              -- )
              return interactiveGoals
          let goalsAtPos : Array InteractiveGoalWithHints := ⟨goalAtPos.foldl (· ++ ·) []⟩

          -- diags are labeled in Lsp-positions, which differ from the lean-internal
          -- positions by `1`.
          let lspPosAt := text.utf8PosToLspPos pos

          let mut diagsAtPos : Array InteractiveDiagnostic :=
            -- `+1` for getting the errors after the line.
            diag.filter (·.range.start.line + 1 == lspPosAt.line)

          diagsAtPos ← addCompletionDiagnostics diagsAtPos goalsAtPos lspPosAt intermediateGoalCount

          intermediateGoalCount := goalsAtPos.size

          steps := steps.push ⟨goalsAtPos, source, diagsAtPos, lspPosAt.line, lspPosAt.character⟩
        else
          -- No goals present
          steps := steps.push default



      --       //   if (goals.length && goalCount > goals.length) {
      -- //     messages.unshift({
      -- //       range: {
      -- //         start: {
      -- //           line: i-1,
      -- //           character: 0,
      -- //         },
      -- //         end: {
      -- //           line: i-1,
      -- //           character: 0,
      -- //         }},
      -- //       severity: DiagnosticSeverity.Information,
      -- //       message: {
      -- //       text: 'intermediate goal solved 🎉'
      -- //       }
      -- //     })
      -- //   }

      -- Level is completed if there are no errrors or warnings

      let completedWithWarnings : Bool := ¬ diag.any (·.severity? == some .error)
      let completed : Bool := completedWithWarnings ∧ ¬ diag.any (·.severity? == some .warning)

      -- Filter out the "unsolved goals" message
      diag := diag.filter (fun d => match d.message with
        | .append ⟨(.text x) :: _⟩ => x != "unsolved goals"
        | _ => true)

      let lastPos := text.utf8PosToLspPos text.positions.back
      let remainingDiags : Array InteractiveDiagnostic :=
        diag.filter (fun d => d.range.start.line >= lastPos.line)

      return some {
        steps := steps,
        diagnostics := remainingDiags,
        completed := completed,
        completedWithWarnings := completedWithWarnings,
        lastPos := lastPos.line
        }

open RequestM in

def getInteractiveGoals (p : Lsp.PlainGoalParams) : RequestM (RequestTask (Option <| InteractiveGoals)) := do
  let doc ← readDoc
  let rc ← readThe RequestContext
  let text := doc.meta.text
  let hoverPos := text.lspPosToUtf8Pos p.position
  -- TODO: I couldn't find a good condition to find the correct snap. So we are looking
  -- for the first snap with goals here:
  withWaitFindSnap doc (fun s => ¬ (s.infoTree.goalsAt? doc.meta.text hoverPos).isEmpty)
    (notFoundX := return none) fun snap => do
      if let rs@(_ :: _) := snap.infoTree.goalsAt? doc.meta.text hoverPos then
        let goals : List <| Array InteractiveGoal ← rs.mapM fun { ctxInfo := ci, tacticInfo := ti, useAfter := useAfter, .. } => do
          let ciAfter := { ci with mctx := ti.mctxAfter }
          let ci := if useAfter then ciAfter else { ci with mctx := ti.mctxBefore }
          -- compute the interactive goals
          let goals ← ci.runMetaM {} do
            return List.toArray <| if useAfter then ti.goalsAfter else ti.goalsBefore
          let goals ← ci.runMetaM {} do
             goals.mapM fun goal => do
              -- let hints ← findHints goal doc.meta rc.initParams
              return ← goalToInteractive goal
          -- compute the goal diff
          -- let goals ← ciAfter.runMetaM {} (do
          --     try
          --       Widget.diffInteractiveGoals useAfter ti goals
          --     catch _ =>
          --       -- fail silently, since this is just a bonus feature
          --       return goals
          -- )
          return goals
        return some <| ⟨goals.foldl (· ++ ·) #[]⟩
      else
        return none

builtin_initialize
  registerBuiltinRpcProcedure
    `Game.getInteractiveGoals
    Lsp.PlainGoalParams
    (Option <| InteractiveGoals
    )
    getInteractiveGoals

builtin_initialize
  registerBuiltinRpcProcedure
    `Game.getProofState
    Lsp.PlainGoalParams
    (Option ProofState)
    getProofState


end GameServer