nornir 0.5.4

//! 💬 **Chat** tab — a "local Claude" / agent-agnostic chat panel wired to the
//! ACTIVE local model (EPIC H / #45).
//!
//! This is the real interactive surface for nornir's local-model story: type a
//! prompt, hit Send, and the message is completed by an in-process
//! [`Generator`](crate::warehouse::generator::Generator) (candle by default,
//! loading a quantized GGUF from the `/mnt/work4t` model cache — see
//! [`candle::models_root`](crate::warehouse::generator::candle::models_root)). No
//! ollama daemon, no cloud API: the completion runs on this machine.
//!
//! ## Active model
//! The model is a `<backend>:<model>` spec ([`ChatTabState::spec`]). It defaults
//! to [`DEFAULT_MODEL_SPEC`] but honours `$NORNIR_GEN_BACKEND` (the CLI-parity
//! `--gen` fallback) so the operator — or the robot harness — points the chat at
//! any backend/model without a rebuild. The generator is built lazily + cached on
//! the first Send (constructing it is cheap; the multi-GB weight load happens on
//! the first `complete`).
//!
//! ## MCP wiring (agent-agnostic)
//! nornir's MCP server (`nornir-mcp`) is a **stdio relay** — an agent spawns it
//! locally and it relays to a `nornir-server` over gRPC. There is no network MCP
//! endpoint to register. [`mcp_client_config`] emits the EXACT `mcpServers` block
//! (command + env) a local agent (Claude Code / Copilot / any MCP client) uses,
//! and it rides in [`ChatTabState::state_json`] so the config is discoverable, and
//! is written to a file at app start (`<warehouse>/nornir-mcp.json`).
//!
//! ## Robot / introspection
//! [`state_json`](ChatTabState::state_json) exposes the transcript + `last_response`
//! so the robot-UI harness reads back the model's reply as DATA (LAW #6). The
//! live-drive keys `chat.input`/`chat.model`/`chat.system`/`chat.max_tokens` (set)
//! and `chat.send`/`chat.clear` (click) run the SAME code the GUI buttons do.

use eframe::egui::{self, RichText, ScrollArea};

use super::facett_theme::{Theme, RED};
use crate::warehouse::generator::{generator, GenRequest, Generator};

/// The default active model: the RECENT, 16 GB-class local model (Qwen2.5-7B
/// Instruct, Q4_K_M — ~4.7 GB, fits an RTX 5080's ~14 GB free VRAM). Overridable
/// via `$NORNIR_GEN_BACKEND` (CLI-parity with `--gen`).
pub const DEFAULT_MODEL_SPEC: &str = "candle:qwen2.5-7b";

/// Token cap per reply (kept modest so a CPU decode stays interactive).
const DEFAULT_MAX_TOKENS: usize = 256;

/// One transcript turn.
struct ChatMsg {
    /// `"user"` or `"assistant"`.
    role: &'static str,
    text: String,
}

/// The 💬 Chat pane state: the active model spec, the transcript, the in-flight
/// input, and the lazily-built generator (cached across sends).
pub struct ChatTabState {
    /// The active `<backend>:<model>` spec (e.g. `candle:qwen2.5-7b`).
    spec: String,
    /// The paste-in prompt buffer.
    input: String,
    /// Optional system prompt sent with every turn.
    system: String,
    /// Hard cap on generated tokens per reply.
    max_tokens: usize,
    /// The full transcript (user + assistant turns).
    messages: Vec<ChatMsg>,
    /// The most recent assistant reply (what the robot asserts on).
    last_response: Option<String>,
    /// The last error, if the model failed to load / complete.
    error: Option<String>,
    /// The lazily-built + cached generator. `None` until the first send; then
    /// `Some(Ok)` (built) or `Some(Err)` (bad spec / not compiled).
    r#gen: Option<Result<Box<dyn Generator>, String>>,
    theme: Theme,
}

impl Default for ChatTabState {
    fn default() -> Self {
        Self::new()
    }
}

impl ChatTabState {
    /// Build the chat pane with the active model resolved from
    /// `$NORNIR_GEN_BACKEND` (CLI-parity) or [`DEFAULT_MODEL_SPEC`]. Cheap — no
    /// weights are loaded until the first Send.
    pub fn new() -> Self {
        let spec = crate::warehouse::generator::spec_from_env()
            .unwrap_or_else(|| DEFAULT_MODEL_SPEC.to_string());
        Self::with_spec(spec)
    }

    /// Build the chat pane pinned to an explicit model spec (tests / callers that
    /// want a specific backend).
    pub fn with_spec(spec: impl Into<String>) -> Self {
        Self {
            spec: spec.into(),
            input: String::new(),
            system: String::new(),
            max_tokens: DEFAULT_MAX_TOKENS,
            messages: Vec::new(),
            last_response: None,
            error: None,
            r#gen: None,
            theme: Theme::default(),
        }
    }

    /// Set the facett palette the pane paints with.
    pub fn set_palette(&mut self, t: Theme) {
        self.theme = t;
    }

    /// The active model spec (`<backend>:<model>`).
    pub fn spec(&self) -> &str {
        &self.spec
    }

    /// Complete the current [`input`](Self::input) with the active model,
    /// appending the user turn + the assistant reply to the transcript. Runs the
    /// generation SYNCHRONOUSLY (like the Funnel `run_demo` drive path) so the
    /// robot harness reads a deterministic `last_response` right after. Returns the
    /// reply text on success.
    pub fn send(&mut self) -> Result<String, String> {
        let prompt = self.input.trim().to_string();
        if prompt.is_empty() {
            return Err("empty prompt (nothing to send)".into());
        }
        // Build + cache the generator on first use (cheap: no weights loaded here).
        if self.r#gen.is_none() {
            self.r#gen = Some(generator(&self.spec).map_err(|e| e.to_string()));
        }
        // Run the completion; the borrow of `gen` ends with this match so we can
        // then mutate the transcript.
        let result: Result<String, String> = match self.r#gen.as_ref().expect("set above") {
            Err(e) => Err(e.clone()),
            Ok(g) => {
                let mut req = GenRequest::new(prompt.clone()).with_max_tokens(self.max_tokens);
                if !self.system.trim().is_empty() {
                    req = req.with_system(self.system.clone());
                }
                g.complete(&req)
                    .map(|a| a.text.trim().to_string())
                    .map_err(|e| format!("{e:#}"))
            }
        };
        self.messages.push(ChatMsg { role: "user", text: prompt });
        match result {
            Ok(text) => {
                self.messages.push(ChatMsg { role: "assistant", text: text.clone() });
                self.last_response = Some(text.clone());
                self.error = None;
                self.input.clear();
                Ok(text)
            }
            Err(e) => {
                self.error = Some(e.clone());
                Err(e)
            }
        }
    }

    /// **R6** drive-half: set a named chat field on the LIVE viz (CLI/robot parity
    /// with `RobotSession::set_field`). Recognised fields:
    ///
    /// * `input` — the paste-in prompt buffer (text);
    /// * `system` — the system prompt (text);
    /// * `model` — the active `<backend>:<model>` spec (rebuilds the generator);
    /// * `max_tokens` — the per-reply token cap (positive integer).
    pub(crate) fn set_field_for_drive(&mut self, field: &str, value: &str) -> Result<(), String> {
        match field {
            "input" => {
                self.input = value.to_string();
                Ok(())
            }
            "system" => {
                self.system = value.to_string();
                Ok(())
            }
            "model" => {
                self.spec = value.trim().to_string();
                // Drop the cached generator so the next send rebuilds for the new spec.
                self.r#gen = None;
                Ok(())
            }
            "max_tokens" => {
                let n: usize = value
                    .trim()
                    .parse()
                    .map_err(|_| format!("max_tokens must be a positive integer, got {value:?}"))?;
                self.max_tokens = n.clamp(1, 4096);
                Ok(())
            }
            other => Err(format!("unknown chat field {other:?}")),
        }
    }

    /// **R6** drive-half: click a named chat button on the LIVE viz (parity with
    /// `RobotSession::click_by_id`). Recognised buttons:
    ///
    /// * `send` — complete the current input with the active model;
    /// * `clear` — reset the transcript.
    pub(crate) fn click_for_drive(
        &mut self,
        button: &str,
        log: &super::action_log::ActionLog,
    ) -> Result<(), String> {
        match button {
            "send" => {
                let r = self.send();
                match &r {
                    Ok(text) => log.push(
                        super::action_log::Kind::Click,
                        format!("chat send (drive): {} char reply", text.len()),
                    ),
                    Err(e) => log.push(
                        super::action_log::Kind::Error,
                        format!("chat send (drive) failed: {e}"),
                    ),
                }
                r.map(|_| ())
            }
            "clear" => {
                self.messages.clear();
                self.last_response = None;
                self.error = None;
                Ok(())
            }
            other => Err(format!("unknown chat button {other:?}")),
        }
    }

    /// The active model's introspection block (which on-disk file backs it, its
    /// footprint, availability, and CPU/GPU device). Candle-aware; other backends
    /// fall back to a probe of id + availability.
    fn model_detail(&self) -> serde_json::Value {
        #[cfg(feature = "gen-candle")]
        if let Some(model) = self.spec.strip_prefix("candle:") {
            return crate::warehouse::generator::candle::spec_detail(model);
        }
        match generator(&self.spec) {
            Ok(g) => serde_json::json!({ "id": g.id(), "available": g.available() }),
            Err(e) => serde_json::json!({ "id": self.spec, "error": e.to_string() }),
        }
    }

    /// 💬 Chat's slice of `state_json` (LAW #6): the active model, the transcript,
    /// the `last_response` the robot asserts on, and the MCP-to-nornir client
    /// config (so an agent's wiring is discoverable from the UI dump).
    pub fn state_json(&self) -> serde_json::Value {
        let messages: Vec<serde_json::Value> = self
            .messages
            .iter()
            .map(|m| serde_json::json!({ "role": m.role, "text": m.text }))
            .collect();
        let turn_count = self.messages.iter().filter(|m| m.role == "user").count();
        serde_json::json!({
            "model_spec": self.spec,
            "model": self.model_detail(),
            "system": self.system,
            "max_tokens": self.max_tokens,
            "input": self.input,
            "turn_count": turn_count,
            "message_count": self.messages.len(),
            "messages": messages,
            "last_response": self.last_response,
            // The single flag the robot chat test asserts: a non-empty reply came back.
            "responded": self.last_response.as_deref().map(|s| !s.is_empty()).unwrap_or(false),
            "error": self.error,
            "mcp": mcp_state_json(),
            "palette": self.theme.name,
        })
    }

    pub fn draw(&mut self, ui: &mut egui::Ui, log: &super::action_log::ActionLog) {
        let theme = self.theme;
        ui.horizontal(|ui| {
            ui.heading("💬 Local chat");
            ui.separator();
            ui.label(RichText::new(&self.spec).monospace().color(theme.accent));
        });
        ui.label(
            "an agent-agnostic chat wired to the ACTIVE local model — a quantized GGUF loaded \
             in-process (candle), cached under /mnt/work4t. No ollama, no cloud.",
        );

        // ── MCP-to-nornir config (discoverable) ─────────────────────────────────
        egui::CollapsingHeader::new("🔌 MCP → nornir (agent wiring)").show(ui, |ui| {
            ui.label(
                "nornir-mcp is a STDIO relay (no network endpoint): an agent spawns it locally; \
                 it relays to nornir-server over gRPC. Wire it in with:",
            );
            ui.monospace(mcp_add_command());
            ui.add_space(4.0);
            ui.label("or the mcpServers block:");
            ui.monospace(
                serde_json::to_string_pretty(&mcp_client_config()).unwrap_or_default(),
            );
        });
        ui.separator();

        // ── transcript ──────────────────────────────────────────────────────────
        ScrollArea::vertical()
            .auto_shrink([false, false])
            .max_height((ui.available_height() - 90.0).max(60.0))
            .stick_to_bottom(true)
            .show(ui, |ui| {
                if self.messages.is_empty() {
                    ui.weak("(no messages yet — type a prompt below and hit Send)");
                }
                for m in &self.messages {
                    let (who, col) = if m.role == "user" {
                        ("you", theme.text)
                    } else {
                        ("model", theme.accent)
                    };
                    ui.horizontal_wrapped(|ui| {
                        ui.label(RichText::new(format!("{who}: ")).strong().color(col));
                        ui.label(&m.text);
                    });
                    ui.add_space(4.0);
                }
                if let Some(err) = &self.error {
                    ui.colored_label(RED, format!("error: {err}"));
                }
            });

        ui.separator();
        // ── input row ─────────────────────────────────────────────────────────────
        ui.horizontal(|ui| {
            let hint = "ask the local model…";
            let edit = egui::TextEdit::singleline(&mut self.input)
                .hint_text(hint)
                .desired_width(ui.available_width() - 90.0);
            let resp = ui.add(edit);
            let submit = (resp.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)))
                || ui.button("➤ Send").clicked();
            if submit && !self.input.trim().is_empty() {
                let _ = self.click_for_drive("send", log);
            }
        });
    }
}

// ─── MCP-to-nornir client config (agent-agnostic, discoverable) ──────────────

/// The nornir-server gRPC target an agent's `nornir-mcp` relay points at
/// (`$NORNIR_SERVER`, default `http://localhost:7878`).
fn mcp_server_target() -> String {
    std::env::var("NORNIR_SERVER").unwrap_or_else(|_| "http://localhost:7878".to_string())
}

/// The EXACT `mcpServers` config block a local agent (Claude Code / Copilot / any
/// agent-agnostic MCP client) uses to reach nornir's tools. nornir-mcp is a STDIO
/// relay — the agent spawns it locally; it relays to `nornir-server` over gRPC. So
/// the config is a spawn command + env, NOT a URL.
pub fn mcp_client_config() -> serde_json::Value {
    let token = std::env::var("NORNIR_SERVER_TOKEN").unwrap_or_else(|_| "<token>".to_string());
    serde_json::json!({
        "mcpServers": {
            "nornir": {
                "command": "nornir-mcp",
                "args": [],
                "env": {
                    "NORNIR_SERVER": mcp_server_target(),
                    "NORNIR_SERVER_TOKEN": token,
                }
            }
        }
    })
}

/// The one-liner `claude mcp add` form of [`mcp_client_config`].
pub fn mcp_add_command() -> String {
    format!(
        "claude mcp add nornir --env NORNIR_SERVER={} \
         --env NORNIR_SERVER_TOKEN=<token> -- nornir-mcp",
        mcp_server_target()
    )
}

/// The MCP block for the chat pane's `state_json` (transport + command + the full
/// client config + the add-command), so an agent's wiring is discoverable from the
/// headless UI dump (LAW #6).
pub fn mcp_state_json() -> serde_json::Value {
    serde_json::json!({
        "transport": "stdio",
        "command": "nornir-mcp",
        "server": mcp_server_target(),
        "config": mcp_client_config(),
        "add_command": mcp_add_command(),
        "note": "nornir-mcp is a local stdio relay; it is NOT a network MCP endpoint. \
                 Spawn it beside the agent; it relays to nornir-server over gRPC.",
    })
}

/// Write the MCP client config to a discoverable file (`<dir>/nornir-mcp.json`) so
/// an agent operator can point their MCP client at nornir without hunting. Called
/// at app start with the warehouse root. Best-effort (a write failure is logged by
/// the caller, never fatal).
pub fn write_mcp_config_file(dir: &std::path::Path) -> std::io::Result<std::path::PathBuf> {
    let path = dir.join("nornir-mcp.json");
    let body = serde_json::to_string_pretty(&mcp_client_config()).unwrap_or_default();
    std::fs::write(&path, body)?;
    Ok(path)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn mcp_config_is_stdio_relay_not_a_url() {
        let cfg = mcp_client_config();
        // The agent spawns nornir-mcp locally (a command), not a URL endpoint.
        assert_eq!(cfg["mcpServers"]["nornir"]["command"], "nornir-mcp");
        assert!(
            cfg["mcpServers"]["nornir"]["env"]["NORNIR_SERVER"]
                .as_str()
                .unwrap()
                .starts_with("http"),
            "relay points at the nornir-server gRPC target"
        );
        assert!(mcp_add_command().contains("claude mcp add nornir"));
        assert_eq!(mcp_state_json()["transport"], "stdio");
    }

    #[test]
    fn empty_input_send_errors_not_panics() {
        let mut chat = ChatTabState::with_spec("mock");
        let err = chat.send().unwrap_err();
        assert!(err.contains("empty prompt"), "{err}");
        assert!(!chat.state_json()["responded"].as_bool().unwrap());
    }

    #[test]
    fn mock_backend_round_trips_a_reply_into_state_json() {
        // Backend-agnostic proof of the whole insert→send→read-back path using the
        // always-available `mock` generator (no weights): set the input via the R6
        // drive verb, click send, and assert the transcript + last_response.
        let log = super::super::action_log::ActionLog::new();
        let mut chat = ChatTabState::with_spec("mock");
        chat.set_field_for_drive("input", "capital of France?").unwrap();
        chat.click_for_drive("send", &log).unwrap();

        let st = chat.state_json();
        assert_eq!(st["turn_count"], 1);
        assert_eq!(st["message_count"], 2, "one user + one assistant turn");
        assert!(st["responded"].as_bool().unwrap(), "a non-empty reply came back");
        let reply = st["last_response"].as_str().unwrap();
        assert!(reply.contains("capital of France?"), "mock echoes the prompt: {reply}");
        // The input was consumed on send.
        assert_eq!(st["input"], "");
    }

    #[test]
    fn set_model_field_rebuilds_generator() {
        let mut chat = ChatTabState::with_spec("mock:one");
        chat.set_field_for_drive("model", "mock:two").unwrap();
        assert_eq!(chat.spec(), "mock:two");
    }
}