/-
  ComonadCommands.lean — Pantograph API handlers for proof navigation.

  Provides the following commands:
    #comonadNode      - Get single node details (node/1 schema)
    #comonadSubgraph  - Get rooted subgraph (subgraph/1 schema)
    #comonadClusters  - Get cluster membership for a node (clusters/1 schema)

  These commands are designed for the Pantograph REPL interface,
  emitting JSON that can be consumed by any client (Mathematica, Python, etc.).
-/

import Infoductor.Comonad.ComonadFinder

open Lean Meta Elab Command

-- ════════════════════════════════════════════════════════════════════════════
-- § 1  JSON serialization for node/1 schema
-- ════════════════════════════════════════════════════════════════════════════

/-- Serialize a ProofNode with cluster membership to node/1 JSON. -/
private def nodeToNodeJson (n : ProofNode) (clusters : Array Nat) : Lean.Json :=
  .mkObj [ ("schema",    .str "node/1")
         , ("contentId", .str n.contentId)
         , ("label",     .str n.label)
         , ("status",    .str n.status)
         , ("shapeKey",  .str n.shapeKey)
         , ("depIds",    .arr (n.depIds.map .str))
         , ("clusters",  .arr (clusters.map fun id => .num ⟨Int.ofNat id, 0⟩)) ]

-- ════════════════════════════════════════════════════════════════════════════
-- § 2  JSON serialization for subgraph/1 schema
-- ════════════════════════════════════════════════════════════════════════════

/-- Serialize an edge to JSON. -/
private def edgeToJson (edge : ContentHash × ContentHash) : Lean.Json :=
  .mkObj [ ("from", .str edge.1)
         , ("to",   .str edge.2) ]

/-- Serialize a subgraph to subgraph/1 JSON. -/
private def subgraphToJson
    (rootId : ContentHash)
    (nodes  : Array ProofNode)
    (edges  : Array (ContentHash × ContentHash))
    (graph  : ProofGraph) : Lean.Json :=
  let nodeJsons := nodes.map fun n =>
    let clusters := findClustersContaining graph n.contentId
    nodeToNodeJson n clusters
  .mkObj [ ("schema", .str "subgraph/1")
         , ("rootId", .str rootId)
         , ("nodes",  .arr nodeJsons)
         , ("edges",  .arr (edges.map edgeToJson)) ]

-- ════════════════════════════════════════════════════════════════════════════
-- § 3  JSON serialization for clusters/1 schema
-- ════════════════════════════════════════════════════════════════════════════

/-- Serialize cluster membership info to clusters/1 JSON. -/
private def clustersToJson
    (nodeId     : ContentHash)
    (clusterIds : Array Nat)
    (graph      : ProofGraph) : Lean.Json :=
  let clusterDetails := clusterIds.filterMap fun cid =>
    graph.clusters.find? (·.id == cid) |>.map fun c =>
      .mkObj [ ("id",          .num ⟨c.id, 0⟩)
             , ("shapeKey",    .str c.shapeKey)
             , ("size",        .num ⟨c.size, 0⟩)
             , ("extractable", .bool c.extractable)
             , ("instanceCount", .num ⟨c.instances.size, 0⟩) ]
  .mkObj [ ("schema",     .str "clusters/1")
         , ("nodeId",     .str nodeId)
         , ("clusterIds", .arr (clusterIds.map fun id => .num ⟨Int.ofNat id, 0⟩))
         , ("clusters",   .arr clusterDetails) ]

-- ════════════════════════════════════════════════════════════════════════════
-- § 4  Commands
-- ════════════════════════════════════════════════════════════════════════════

/-- `#comonadNode "nodeId" from theoremName`

    Get details of a single node in the proof graph.
    Emits node/1 JSON with cluster membership. -/
elab "#comonadNode " nodeId:str " from " thm:ident : command => do
  let name := thm.getId
  let nid  := nodeId.getString
  let graph ← liftTermElabM (findComonadsCore name)
  match getNodeInfo graph nid with
  | none   => logError m!"Node '{nid}' not found in proof of {name}"
  | some n =>
    let clusters := findClustersContaining graph nid
    logInfo m!"{(nodeToNodeJson n clusters).compress}"

/-- `#comonadSubgraph "rootId" from theoremName`

    Get the subgraph rooted at a specific node.
    Emits subgraph/1 JSON with nodes and edges. -/
elab "#comonadSubgraph " rootId:str " from " thm:ident : command => do
  let name := thm.getId
  let rid  := rootId.getString
  let graph ← liftTermElabM (findComonadsCore name)
  match getNodeInfo graph rid with
  | none   => logError m!"Root node '{rid}' not found in proof of {name}"
  | some _ =>
    let nodes := getSubgraph graph rid
    let edges := getSubgraphEdges graph rid
    logInfo m!"{(subgraphToJson rid nodes edges graph).compress}"

/-- `#comonadClusters "nodeId" from theoremName`

    Find all clusters that contain a given node.
    Emits clusters/1 JSON with cluster details. -/
elab "#comonadClusters " nodeId:str " from " thm:ident : command => do
  let name := thm.getId
  let nid  := nodeId.getString
  let graph ← liftTermElabM (findComonadsCore name)
  let clusterIds := findClustersContaining graph nid
  logInfo m!"{(clustersToJson nid clusterIds graph).compress}"

-- (Test section dropped on port — the original mm-lean tests
-- referenced theorems defined in ComonadFinder's Test section
-- that we also dropped.  Restore alongside a Test/ harness that
-- handles the Lean 4 v4.30 theorem-value-access change.)