//! Embedded MCP server. Lets a Claude session running anywhere on the
//! same machine — including inside one of tiletopia's own panes — inspect
//! the workspace via Model Context Protocol.
//!
//! V1 surface (read-only):
//!   resources: tiletopia://layout, tiletopia://panes, tiletopia://hosts
//!   tools:     read_pane(leaf_id, last_lines?, after_seq?)
//!              wait_for_idle(leaf_id, idle_ms?, timeout_ms?)
//!
//! Per-pane `mcpAllow` gate (default-deny) lives in the frontend tree;
//! the frontend mirrors the gated subset into {@link McpState} via the
//! `mcp_update_state` Tauri command. The MCP server only sees what the
//! mirror exposes — no peeking around it.

use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use anyhow::{Context, Result};
use axum::{
    body::Body,
    http::{HeaderMap, HeaderValue, Request, StatusCode},
    middleware::{self, Next},
    response::Response,
    Router,
};
use parking_lot::Mutex as PlMutex;
use rmcp::{
    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
    model::*,
    schemars, tool, tool_handler, tool_router,
    service::RequestContext,
    transport::streamable_http_server::{
        session::local::LocalSessionManager,
        tower::{StreamableHttpServerConfig, StreamableHttpService},
    },
    ErrorData as McpError, RoleServer, ServerHandler,
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tauri::{AppHandle, Emitter, Manager};
use tokio::{net::TcpListener, sync::RwLock, task::JoinHandle};
use tokio_util::sync::CancellationToken;

use crate::pty::{PaneId, PtyManager};

/// Default port for the MCP server. Picked from the IANA-unassigned
/// 47000-range so it's unlikely to collide with anything else on a dev box.
/// Override by editing `port` in `%APPDATA%\com.megaproxy.tiletopia\mcp.json`.
pub const DEFAULT_PORT: u16 = 47821;

const MCP_CONFIG_FILE: &str = "mcp.json";

/// Persisted across restarts so the firewall rule + Claude config snippet
/// don't need re-pasting every launch. Generated on first start.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpPersistedConfig {
    pub port: u16,
    pub token: String,
}

impl McpPersistedConfig {
    fn new_default() -> Self {
        Self {
            port: DEFAULT_PORT,
            token: generate_token(),
        }
    }
}

fn generate_token() -> String {
    use rand::RngCore;
    let mut buf = [0u8; 32];
    rand::rng().fill_bytes(&mut buf);
    hex::encode(buf)
}

fn config_path(app: &AppHandle) -> Result<PathBuf> {
    let dir = app
        .path()
        .app_config_dir()
        .map_err(|e| anyhow::anyhow!("app_config_dir: {e}"))?;
    Ok(dir.join(MCP_CONFIG_FILE))
}

/// Load saved config, or generate-and-save a fresh one on first run.
pub fn load_or_init_config(app: &AppHandle) -> Result<McpPersistedConfig> {
    let path = config_path(app)?;
    if path.exists() {
        let raw = std::fs::read_to_string(&path).context("read mcp.json")?;
        let cfg: McpPersistedConfig =
            serde_json::from_str(&raw).context("parse mcp.json")?;
        return Ok(cfg);
    }
    let cfg = McpPersistedConfig::new_default();
    save_config(app, &cfg)?;
    Ok(cfg)
}

pub fn save_config(app: &AppHandle, cfg: &McpPersistedConfig) -> Result<()> {
    let path = config_path(app)?;
    if let Some(dir) = path.parent() {
        std::fs::create_dir_all(dir).context("create_dir_all")?;
    }
    let tmp = path.with_extension("json.tmp");
    let json = serde_json::to_string_pretty(cfg).context("serialize mcp cfg")?;
    std::fs::write(&tmp, json.as_bytes()).context("write tmp mcp.json")?;
    // Atomic on Unix; MoveFileEx with REPLACE_EXISTING on Windows.
    std::fs::rename(&tmp, &path).context("rename mcp.json")?;
    Ok(())
}

// ----------------------------------------------------------------------------
// Shared state mirrored from the frontend.
// ----------------------------------------------------------------------------

pub type LeafId = String;

/// Cached snapshot the frontend pushes via `mcp_update_state` whenever the
/// tree or hosts change. Source of truth for everything except scrollback,
/// which the backend collects directly via {@link PtyManager}.
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpMirror {
    /// Serialised layout tree (full structure, post-filtering happens
    /// per-resource — see read_resource).
    #[serde(default)]
    pub layout_json: String,
    /// Map of leaf id → pane metadata. Includes only leaves with
    /// `mcpAllow === true` (frontend gates before mirroring).
    #[serde(default)]
    pub leaves: HashMap<LeafId, MirroredLeaf>,
    /// Saved SSH hosts, password fields stripped.
    #[serde(default)]
    pub hosts: Vec<MirroredHost>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MirroredLeaf {
    pub pane_id: Option<PaneId>,
    pub label: Option<String>,
    pub shell_kind: String,
    pub distro: Option<String>,
    pub ssh_host_id: Option<String>,
    #[serde(default)]
    pub broadcast: bool,
    #[serde(default)]
    pub active: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MirroredHost {
    pub id: String,
    pub label: String,
    pub hostname: String,
    pub user: Option<String>,
    pub port: Option<u16>,
    #[serde(default)]
    pub has_password: bool,
}

#[derive(Default)]
pub struct McpState {
    pub bearer_token: String,
    pub mirror: McpMirror,
}

// ----------------------------------------------------------------------------
// Action reply registry.
// ----------------------------------------------------------------------------

/// Registry of pending frontend action requests. Each entry maps a `requestId`
/// to a oneshot sender that the `mcp_action_reply` Tauri command will fire
/// once the frontend resolves or rejects the action.
///
/// Owned as separate managed state (Arc<PendingActions>) so Tauri commands can
/// grab it via `tauri::State<'_, Arc<PendingActions>>` without needing to lock
/// the entire McpState or pass TileService around.
pub struct PendingActions(
    pub PlMutex<HashMap<String, tokio::sync::oneshot::Sender<Result<serde_json::Value, String>>>>,
);

impl Default for PendingActions {
    fn default() -> Self {
        Self(PlMutex::new(HashMap::new()))
    }
}

// ----------------------------------------------------------------------------
// Audit / request event payload types.
// ----------------------------------------------------------------------------

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct McpActionRequest {
    request_id: String,
    tool: &'static str,
    args: serde_json::Value,
    needs_confirm: bool,
    reason: Option<String>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase", tag = "kind")]
enum McpAuditResult {
    Ok,
    Denied { reason: String, hard: bool },
    Failed { msg: String },
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct McpAuditEntry {
    ts_ms: u64,
    tool: &'static str,
    args_summary: String,
    result: McpAuditResult,
    duration_ms: u64,
}

fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

fn truncate_summary(s: &str) -> String {
    if s.len() > 80 {
        format!("{}...", &s[..80])
    } else {
        s.to_string()
    }
}

// ----------------------------------------------------------------------------
// MCP service: tools + resources.
// ----------------------------------------------------------------------------

#[derive(Clone)]
pub struct TileService {
    ptys: Arc<PtyManager>,
    state: Arc<RwLock<McpState>>,
    pending: Arc<PendingActions>,
    app: AppHandle,
    tool_router: ToolRouter<Self>,
}

#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct ReadPaneArgs {
    /// Stable leaf id from the tree (uuid-shaped). Must belong to a pane
    /// the user has allow-listed for MCP access.
    pub leaf_id: LeafId,
    /// Return only the last N lines (default 200, hard cap 3000).
    #[serde(default)]
    pub last_lines: Option<usize>,
    /// Only return bytes whose seq > this. Pair with the `__seq__` value
    /// returned in a prior call for incremental polling.
    #[serde(default)]
    pub after_seq: Option<u64>,
}

#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct WaitForIdleArgs {
    pub leaf_id: LeafId,
    /// Required quiet window before declaring idle (default 500 ms).
    #[serde(default)]
    pub idle_ms: Option<u64>,
    /// Hard timeout in ms; returns timeout=true after this (default 30s,
    /// hard cap 5 min).
    #[serde(default)]
    pub timeout_ms: Option<u64>,
}

#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct SetLabelArgs {
    /// Stable leaf id from the tree (uuid-shaped). Must belong to a pane
    /// the user has allow-listed for MCP access.
    pub leaf_id: LeafId,
    /// New human-readable label. Pass an empty string to clear the label.
    pub label: String,
}

const READ_PANE_HARD_CAP_LINES: usize = 3000;
const WAIT_TIMEOUT_HARD_CAP_MS: u64 = 5 * 60 * 1000;

#[tool_router]
impl TileService {
    pub fn new(
        ptys: Arc<PtyManager>,
        state: Arc<RwLock<McpState>>,
        pending: Arc<PendingActions>,
        app: AppHandle,
    ) -> Self {
        Self {
            ptys,
            state,
            pending,
            app,
            tool_router: Self::tool_router(),
        }
    }

    /// Dispatch an action that the frontend must execute. Generates a uuid,
    /// registers a oneshot, emits the "mcp://request" event with the args,
    /// awaits the reply with a 30s timeout. Also emits an "mcp://audit" event
    /// after the call resolves (regardless of outcome).
    async fn dispatch_action(
        &self,
        tool: &'static str,
        args: serde_json::Value,
        args_repr: String,
    ) -> Result<serde_json::Value, McpError> {
        let start_ms = now_ms();
        let args_summary = truncate_summary(&args_repr);

        tracing::debug!(tool, args_repr = %args_repr, "dispatch_action: start");

        // 1. Load user policy.
        let policy = crate::mcp_policy::load_or_init(&self.app)
            .map_err(|e| McpError::internal_error(e.to_string(), None))?;

        // 2. Hard-deny check (for any tool — is_hard_denied checks for shell
        //    catastrophe patterns; for non-write_pane tools the patterns are
        //    unlikely to match args_repr but the check is cheap and safe).
        if let Some(label) = crate::mcp_policy::is_hard_denied(&args_repr) {
            let duration_ms = now_ms() - start_ms;
            let audit = McpAuditEntry {
                ts_ms: start_ms,
                tool,
                args_summary: args_summary.clone(),
                result: McpAuditResult::Denied {
                    reason: label.to_string(),
                    hard: true,
                },
                duration_ms,
            };
            tracing::debug!(tool, reason = label, hard = true, "dispatch_action: hard-denied");
            let _ = self.app.emit("mcp://audit", &audit);
            return Err(McpError::invalid_params(
                format!("hard-denied: {label}"),
                None,
            ));
        }

        // 3. Evaluate user-policy decision.
        let decision = crate::mcp_policy::evaluate(&policy, tool, &args_repr);

        tracing::debug!(tool, ?decision, "dispatch_action: policy decision");

        // 4. Handle Deny.
        let (needs_confirm, ask_reason) = match &decision {
            crate::mcp_policy::PolicyDecision::Allow => (false, None),
            crate::mcp_policy::PolicyDecision::Ask { reason } => {
                (true, Some(reason.clone()))
            }
            crate::mcp_policy::PolicyDecision::Deny { reason, hard } => {
                let duration_ms = now_ms() - start_ms;
                let audit = McpAuditEntry {
                    ts_ms: start_ms,
                    tool,
                    args_summary: args_summary.clone(),
                    result: McpAuditResult::Denied {
                        reason: reason.clone(),
                        hard: *hard,
                    },
                    duration_ms,
                };
                tracing::debug!(tool, reason = %reason, hard, "dispatch_action: denied by policy");
                let _ = self.app.emit("mcp://audit", &audit);
                return Err(McpError::invalid_params(
                    format!("denied: {reason}"),
                    None,
                ));
            }
        };

        // 5. Generate a unique request id, register oneshot, emit mcp://request.
        // uuid crate is not in Cargo.toml; generate via rand (already a dep).
        // TODO: if uuid (v4 feature) is added to Cargo.toml, replace with:
        //   let request_id = uuid::Uuid::new_v4().to_string();
        let request_id = {
            use rand::RngCore;
            let mut bytes = [0u8; 16];
            rand::rng().fill_bytes(&mut bytes);
            // Format as a RFC-4122-style UUID v4 string for frontend interop.
            bytes[6] = (bytes[6] & 0x0f) | 0x40;
            bytes[8] = (bytes[8] & 0x3f) | 0x80;
            format!(
                "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
                bytes[0], bytes[1], bytes[2], bytes[3],
                bytes[4], bytes[5],
                bytes[6], bytes[7],
                bytes[8], bytes[9],
                bytes[10], bytes[11], bytes[12], bytes[13], bytes[14], bytes[15],
            )
        };

        let (tx, rx) = tokio::sync::oneshot::channel();
        {
            self.pending.0.lock().insert(request_id.clone(), tx);
        }

        let payload = McpActionRequest {
            request_id: request_id.clone(),
            tool,
            args,
            needs_confirm,
            reason: ask_reason,
        };
        tracing::debug!(tool, request_id = %request_id, needs_confirm, "dispatch_action: emitting mcp://request");
        let _ = self.app.emit("mcp://request", &payload);

        // 6. Await reply with 30s timeout.
        let result = tokio::time::timeout(Duration::from_secs(30), rx).await;

        let duration_ms = now_ms() - start_ms;

        match result {
            Err(_elapsed) => {
                // Timeout — remove stale sender from registry.
                self.pending.0.lock().remove(&request_id);
                let audit = McpAuditEntry {
                    ts_ms: start_ms,
                    tool,
                    args_summary,
                    result: McpAuditResult::Failed {
                        msg: "timeout".into(),
                    },
                    duration_ms,
                };
                tracing::debug!(tool, request_id = %request_id, "dispatch_action: timed out");
                let _ = self.app.emit("mcp://audit", &audit);
                Err(McpError::internal_error(
                    "action timed out waiting for frontend response",
                    Some(json!({ "requestId": request_id })),
                ))
            }
            Ok(Err(_recv_err)) => {
                // Sender was dropped (shouldn't happen normally).
                let audit = McpAuditEntry {
                    ts_ms: start_ms,
                    tool,
                    args_summary,
                    result: McpAuditResult::Failed {
                        msg: "channel closed".into(),
                    },
                    duration_ms,
                };
                tracing::debug!(tool, request_id = %request_id, "dispatch_action: channel closed");
                let _ = self.app.emit("mcp://audit", &audit);
                Err(McpError::internal_error(
                    "action channel closed unexpectedly",
                    Some(json!({ "requestId": request_id })),
                ))
            }
            Ok(Ok(reply)) => {
                // 7. On reply: emit audit, propagate. Destructure with
                //    ownership so the success payload and the error string
                //    move out cleanly (avoids borrow-then-move on `reply`).
                let (audit_result, err, ok_payload) = match reply {
                    Ok(v) => {
                        tracing::debug!(tool, request_id = %request_id, "dispatch_action: reply ok");
                        (McpAuditResult::Ok, None, Some(v))
                    }
                    Err(msg) => {
                        tracing::debug!(tool, request_id = %request_id, error = %msg, "dispatch_action: reply error");
                        (
                            McpAuditResult::Failed { msg: msg.clone() },
                            Some(McpError::internal_error(msg, None)),
                            None,
                        )
                    }
                };
                let audit = McpAuditEntry {
                    ts_ms: start_ms,
                    tool,
                    args_summary,
                    result: audit_result,
                    duration_ms,
                };
                let _ = self.app.emit("mcp://audit", &audit);
                match err {
                    Some(e) => Err(e),
                    None => Ok(ok_payload.expect("ok branch always sets ok_payload")),
                }
            }
        }
    }

    /// Look up a leaf_id → pane_id under the MCP-allow gate.
    async fn resolve_pane(&self, leaf_id: &str) -> Result<PaneId, McpError> {
        let st = self.state.read().await;
        let leaf = st.mirror.leaves.get(leaf_id).ok_or_else(|| {
            McpError::invalid_params(
                "unknown leaf_id (not visible to MCP; user may need to allow it)",
                Some(json!({ "leaf_id": leaf_id })),
            )
        })?;
        leaf.pane_id.ok_or_else(|| {
            McpError::invalid_params(
                "leaf has no live pane",
                Some(json!({ "leaf_id": leaf_id })),
            )
        })
    }

    #[tool(description = "Read the recent scrollback of a terminal pane. \
        Returns text plus a __seq__=N marker that can be passed back as \
        after_seq for incremental polling.")]
    async fn read_pane(
        &self,
        Parameters(args): Parameters<ReadPaneArgs>,
    ) -> Result<CallToolResult, McpError> {
        let pane_id = self.resolve_pane(&args.leaf_id).await?;

        let ring = self.ptys.ring(pane_id).ok_or_else(|| {
            McpError::internal_error(
                "pane ring missing (pane may have just exited)",
                Some(json!({ "leaf_id": args.leaf_id })),
            )
        })?;
        let (bytes, seq) = {
            let g = ring.lock();
            g.snapshot()
        };

        // Trim by after_seq if provided: bytes in the ring beyond
        // `after_seq` is `seq - after_seq`, clamped against ring size.
        let start = match args.after_seq {
            Some(prev) if seq > prev => {
                let new_bytes = (seq - prev) as usize;
                bytes.len().saturating_sub(new_bytes)
            }
            Some(_) => bytes.len(),
            None => 0,
        };
        let tail = &bytes[start..];

        let text = String::from_utf8_lossy(tail);
        let cap = args
            .last_lines
            .map(|n| n.min(READ_PANE_HARD_CAP_LINES))
            .unwrap_or(200);
        let limited: String = if cap == 0 {
            String::new()
        } else {
            let lines: Vec<&str> = text.lines().collect();
            let start_line = lines.len().saturating_sub(cap);
            lines[start_line..].join("\n")
        };

        Ok(CallToolResult::success(vec![
            Content::text(limited),
            Content::text(format!("__seq__={seq}")),
        ]))
    }

    #[tool(description = "Block until a pane has been quiet (no output) \
        for idle_ms, or timeout_ms elapses. Useful for command-completion \
        synchronisation. Returns {idle:bool, seq:u64, elapsed_ms:u64}.")]
    async fn wait_for_idle(
        &self,
        Parameters(args): Parameters<WaitForIdleArgs>,
    ) -> Result<CallToolResult, McpError> {
        let pane_id = self.resolve_pane(&args.leaf_id).await?;
        let ring = self.ptys.ring(pane_id).ok_or_else(|| {
            McpError::internal_error("pane ring missing", None)
        })?;

        let idle_target = Duration::from_millis(args.idle_ms.unwrap_or(500));
        let timeout = Duration::from_millis(
            args.timeout_ms
                .unwrap_or(30_000)
                .min(WAIT_TIMEOUT_HARD_CAP_MS),
        );
        let start = Instant::now();
        let mut last_seq = ring.lock().snapshot().1;
        let mut last_change = Instant::now();

        loop {
            // Sleep in small slices so we notice both incoming data and
            // the overall timeout promptly.
            tokio::time::sleep(Duration::from_millis(50)).await;
            let now_seq = ring.lock().snapshot().1;
            if now_seq != last_seq {
                last_seq = now_seq;
                last_change = Instant::now();
            }
            if last_change.elapsed() >= idle_target {
                return Ok(CallToolResult::success(vec![Content::text(
                    json!({
                        "idle": true,
                        "seq": last_seq,
                        "elapsed_ms": start.elapsed().as_millis() as u64,
                    })
                    .to_string(),
                )]));
            }
            if start.elapsed() >= timeout {
                return Ok(CallToolResult::success(vec![Content::text(
                    json!({
                        "idle": false,
                        "seq": last_seq,
                        "elapsed_ms": start.elapsed().as_millis() as u64,
                    })
                    .to_string(),
                )]));
            }
        }
    }

    #[tool(description = "Set or clear the human-readable label on a pane. \
        Pass empty string to clear. The leaf must be MCP-allowed.")]
    async fn set_label(
        &self,
        Parameters(args): Parameters<SetLabelArgs>,
    ) -> Result<CallToolResult, McpError> {
        // Validate leaf exists in mirror + is visible to MCP (mcpAllow=true
        // is enforced by the frontend before mirroring, so presence here
        // implies the user has allowed it).
        let _leaf = self
            .state
            .read()
            .await
            .mirror
            .leaves
            .get(&args.leaf_id)
            .cloned()
            .ok_or_else(|| {
                McpError::invalid_params(
                    "unknown leaf_id (not visible to MCP; user may need to allow it)",
                    Some(json!({ "leaf_id": &args.leaf_id })),
                )
            })?;

        let args_repr = format!("leafId={} label={}", &args.leaf_id, &args.label);
        let args_json = json!({ "leafId": &args.leaf_id, "label": &args.label });

        tracing::debug!(leaf_id = %args.leaf_id, label = %args.label, "set_label: dispatching");
        let _ = self
            .dispatch_action("set_label", args_json, args_repr)
            .await?;

        Ok(CallToolResult::success(vec![Content::text("ok")]))
    }
}

#[tool_handler]
impl ServerHandler for TileService {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::new(
            ServerCapabilities::builder()
                .enable_tools()
                .enable_resources()
                .build(),
        )
        .with_server_info(Implementation::from_build_env())
        .with_protocol_version(ProtocolVersion::V_2024_11_05)
        .with_instructions(
            "Tiletopia MCP (read-only v1). Resources: tiletopia://layout, \
             tiletopia://panes, tiletopia://hosts. Tools: read_pane, \
             wait_for_idle. Only panes the user has allow-listed are \
             visible.",
        )
    }

    async fn list_resources(
        &self,
        _r: Option<PaginatedRequestParams>,
        _: RequestContext<RoleServer>,
    ) -> Result<ListResourcesResult, McpError> {
        Ok(ListResourcesResult {
            resources: vec![
                RawResource::new("tiletopia://layout", "layout").no_annotation(),
                RawResource::new("tiletopia://panes", "panes").no_annotation(),
                RawResource::new("tiletopia://hosts", "hosts").no_annotation(),
            ],
            next_cursor: None,
            meta: None,
        })
    }

    async fn read_resource(
        &self,
        req: ReadResourceRequestParams,
        _: RequestContext<RoleServer>,
    ) -> Result<ReadResourceResult, McpError> {
        let state = self.state.read().await;
        let body = match req.uri.as_str() {
            "tiletopia://layout" => state.mirror.layout_json.clone(),
            "tiletopia://panes" => {
                serde_json::to_string(&state.mirror.leaves).unwrap_or_default()
            }
            "tiletopia://hosts" => {
                serde_json::to_string(&state.mirror.hosts).unwrap_or_default()
            }
            other => {
                return Err(McpError::resource_not_found(
                    "resource_not_found",
                    Some(json!({ "uri": other })),
                ));
            }
        };
        Ok(ReadResourceResult::new(vec![ResourceContents::text(
            body, req.uri,
        )]))
    }

    async fn list_resource_templates(
        &self,
        _r: Option<PaginatedRequestParams>,
        _: RequestContext<RoleServer>,
    ) -> Result<ListResourceTemplatesResult, McpError> {
        Ok(ListResourceTemplatesResult {
            resource_templates: vec![],
            next_cursor: None,
            meta: None,
        })
    }
}

// ----------------------------------------------------------------------------
// HTTP wiring + bearer auth.
// ----------------------------------------------------------------------------

async fn bearer_auth(
    axum::extract::State(expected): axum::extract::State<Arc<String>>,
    headers: HeaderMap,
    req: Request<Body>,
    next: Next,
) -> Result<Response, Response> {
    // OAuth-discovery clients probe /.well-known/* and /register before any
    // /mcp request. Letting those fall through to axum's default 404 keeps
    // us out of the OAuth challenge/response game — bearer enforcement only
    // applies to the real MCP surface.
    if !req.uri().path().starts_with("/mcp") {
        return Ok(next.run(req).await);
    }

    let auth_header = headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok());
    let supplied = auth_header.and_then(|s| s.strip_prefix("Bearer "));
    let ok = supplied
        .map(|t| constant_time_eq(t.as_bytes(), expected.as_bytes()))
        .unwrap_or(false);

    tracing::debug!(
        method = %req.method(),
        path = %req.uri().path(),
        auth_present = auth_header.is_some(),
        bearer_present = supplied.is_some(),
        token_match = ok,
        "MCP request"
    );

    if ok {
        return Ok(next.run(req).await);
    }

    let mut resp = Response::builder()
        .status(StatusCode::UNAUTHORIZED)
        .body(Body::empty())
        .unwrap();
    resp.headers_mut().insert(
        axum::http::header::WWW_AUTHENTICATE,
        HeaderValue::from_static(r#"Bearer realm="tiletopia""#),
    );
    Err(resp)
}

fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut d = 0u8;
    for (x, y) in a.iter().zip(b) {
        d |= x ^ y;
    }
    d == 0
}

// ----------------------------------------------------------------------------
// Lifecycle.
// ----------------------------------------------------------------------------

pub struct RunningServer {
    pub addr: SocketAddr,
    pub token: String,
    pub cancel: CancellationToken,
    pub task: JoinHandle<()>,
}

#[derive(Default)]
pub struct McpServerHandle(pub PlMutex<Option<RunningServer>>);

pub async fn start_server(
    app_handle: AppHandle,
    ptys: Arc<PtyManager>,
    state: Arc<RwLock<McpState>>,
    pending: Arc<PendingActions>,
) -> Result<RunningServer> {
    let cfg = load_or_init_config(&app_handle)?;
    let token = cfg.token.clone();
    state.write().await.bearer_token = token.clone();

    let cancel = CancellationToken::new();

    // Fresh service per session; cheap because we share state via Arcs.
    let ptys_f = ptys.clone();
    let state_f = state.clone();
    let pending_f = pending.clone();
    // Clone AppHandle before the move closure so we can pass it into each
    // TileService instance. AppHandle is cheap to clone (it's an Arc inside).
    let app_handle_for_service = app_handle.clone();
    // Disable rmcp's DNS-rebinding host allowlist. The default only permits
    // localhost / 127.0.0.1 / ::1; legitimate WSL clients connect via the
    // dynamic WSL gateway IP (172.x.x.1) which can't be in any static list.
    // Bearer-token auth on /mcp is the real gatekeeper, and we're not
    // running in a browser context where DNS rebinding is a concern.
    let mcp_service = StreamableHttpService::new(
        move || {
            Ok(TileService::new(
                ptys_f.clone(),
                state_f.clone(),
                pending_f.clone(),
                app_handle_for_service.clone(),
            ))
        },
        LocalSessionManager::default().into(),
        StreamableHttpServerConfig::default().disable_allowed_hosts(),
    );

    let router = Router::new()
        .nest_service("/mcp", mcp_service)
        .layer(middleware::from_fn_with_state(
            Arc::new(token.clone()),
            bearer_auth,
        ));

    // Bind to all interfaces so WSL distros in NAT mode can reach the server
    // via the Windows host's WSL-side adapter IP. Auth is bearer-token only.
    // We report 127.0.0.1 in the URL since that's the canonical Windows-side
    // hostname (WSL clients swap in the gateway IP).
    //
    // Try the saved port first so the user's firewall rule + Claude config
    // survive restarts. If it's taken, fall back to an OS-picked port and
    // leave the saved port alone — the conflict may clear later.
    let listener = match TcpListener::bind(("0.0.0.0", cfg.port)).await {
        Ok(l) => l,
        Err(e) => {
            tracing::warn!(
                "MCP saved port {} unavailable ({}); falling back to OS-picked port",
                cfg.port,
                e
            );
            TcpListener::bind("0.0.0.0:0").await?
        }
    };
    let port = listener.local_addr()?.port();
    let addr = SocketAddr::from(([127, 0, 0, 1], port));

    let cancel_inner = cancel.clone();
    let task = tokio::spawn(async move {
        let _ = axum::serve(listener, router)
            .with_graceful_shutdown(async move {
                cancel_inner.cancelled().await;
            })
            .await;
    });

    tracing::info!("MCP server listening on http://{addr}/mcp");
    Ok(RunningServer {
        addr,
        token,
        cancel,
        task,
    })
}

pub fn stop_server(handle: &McpServerHandle) {
    if let Some(srv) = handle.0.lock().take() {
        srv.cancel.cancel();
        srv.task.abort();
        tracing::info!("MCP server stopped");
    }
}

/// Mint a new bearer token, persist it, and return the new value. Caller is
/// responsible for restarting the server if it was running — the live auth
/// middleware captures the token by value at start time.
pub fn regenerate_token(app: &AppHandle) -> Result<String> {
    let mut cfg = load_or_init_config(app)?;
    cfg.token = generate_token();
    save_config(app, &cfg)?;
    Ok(cfg.token)
}