use std::collections::HashMap;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{ChildStdin, Command};
use tokio::sync::{Mutex, mpsc, oneshot};
use tracing::{debug, warn};

use crate::{McpCallResult, McpLogLine, McpLogTx, McpServerClient, McpTool, extract_text, interpolate_env};
use crate::config::McpServerConfig;

/// A server-initiated notification from an MCP server: `(server_name, full JSON-RPC message)`.
pub type McpNotification = (String, Value);

const CALL_TIMEOUT_SECS: u64 = 120;

// ── Elicitation (MCP spec 2025-06-18) ──────────────────────────────────────────

/// A server-initiated elicitation request: the server asks the user for input
/// *during* a tool call (`elicitation/create`). The secret/value never reaches
/// the LLM and is never persisted.
#[derive(Debug, Clone)]
pub struct ElicitationRequest {
    /// Human-readable message to show the user.
    pub message:          String,
    /// JSON Schema (flat object of typed fields) describing the requested input.
    pub requested_schema: Value,
}

/// The user's decision on an elicitation request.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElicitationAction {
    Accept,
    Decline,
    Cancel,
}

impl ElicitationAction {
    fn as_str(self) -> &'static str {
        match self {
            ElicitationAction::Accept  => "accept",
            ElicitationAction::Decline => "decline",
            ElicitationAction::Cancel  => "cancel",
        }
    }
}

/// The reply sent back to the server for an `elicitation/create` request.
#[derive(Debug, Clone)]
pub struct ElicitationReply {
    pub action:  ElicitationAction,
    /// Field values for `accept`; `None` for `decline`/`cancel`.
    pub content: Option<Value>,
}

/// Bridges a server-initiated elicitation to whatever surfaces it to the user
/// (in Skald: the Agent Inbox via `ElicitationManager`). The crate writes the
/// JSON-RPC reply itself — the handler only produces the decision. The returned
/// value (and any secret it carries) flows straight to the server's stdin and is
/// never logged here.
#[async_trait]
pub trait ElicitationHandler: Send + Sync {
    async fn handle(&self, server_name: &str, request: ElicitationRequest) -> ElicitationReply;
}

/// Serialises `msg` as a single JSON-RPC line and writes it to the child's stdin.
async fn write_json_line(stdin: &Arc<Mutex<ChildStdin>>, msg: &Value) {
    if let Ok(mut line) = serde_json::to_string(msg) {
        line.push('\n');
        let _ = stdin.lock().await.write_all(line.as_bytes()).await;
    }
}

/// Handles a server→client JSON-RPC request (e.g. `elicitation/create`) without
/// blocking the read-loop: the user reply may take minutes, so the work is
/// spawned and the JSON-RPC response is written back when it resolves.
fn handle_server_request(
    server_name:          &str,
    msg:                  Value,
    stdin:                &Arc<Mutex<ChildStdin>>,
    handler:              &Option<Arc<dyn ElicitationHandler>>,
    pending_elicitations: &Arc<AtomicUsize>,
) {
    let id     = msg.get("id").cloned().unwrap_or(Value::Null);
    let method = msg.get("method").and_then(Value::as_str).unwrap_or("");

    if method == "elicitation/create" {
        let params  = msg.get("params").cloned().unwrap_or_else(|| json!({}));
        let request = ElicitationRequest {
            message:          params.get("message").and_then(Value::as_str).unwrap_or("").to_string(),
            requested_schema: params.get("requestedSchema").cloned().unwrap_or_else(|| json!({})),
        };
        let stdin = Arc::clone(stdin);

        let Some(handler) = handler.clone() else {
            // Capability declared but no handler wired: cancel so the server
            // doesn't hang waiting for input we can't collect.
            tokio::spawn(async move {
                write_json_line(&stdin, &json!({
                    "jsonrpc": "2.0", "id": id, "result": { "action": "cancel" },
                })).await;
            });
            return;
        };

        pending_elicitations.fetch_add(1, Ordering::SeqCst);
        let counter = Arc::clone(pending_elicitations);
        let server  = server_name.to_string();
        tokio::spawn(async move {
            let reply      = handler.handle(&server, request).await;
            let mut result = json!({ "action": reply.action.as_str() });
            if reply.action == ElicitationAction::Accept {
                if let Some(content) = reply.content {
                    result["content"] = content;
                }
            }
            write_json_line(&stdin, &json!({
                "jsonrpc": "2.0", "id": id, "result": result,
            })).await;
            counter.fetch_sub(1, Ordering::SeqCst);
        });
    } else {
        // Unknown server→client request: reply method-not-found so the server
        // isn't left hanging.
        let stdin  = Arc::clone(stdin);
        let method = method.to_string();
        tokio::spawn(async move {
            write_json_line(&stdin, &json!({
                "jsonrpc": "2.0", "id": id,
                "error": { "code": -32601, "message": format!("method not found: {method}") },
            })).await;
        });
    }
}

/// Guards an in-flight `tools/call`: if this is dropped while still armed — the
/// caller aborted the tool (a `/stop` drops the work future) or the call timed out
/// — it tells the server to stop via `notifications/cancelled` and drops the now
/// orphaned `pending` entry (which would otherwise leak). Disarmed once the server
/// replies or disconnects, where cancelling is pointless. The spec forbids
/// cancelling `initialize`, so only `tools/call` arms a guard.
struct CancelOnDrop {
    id:      u64,
    stdin:   Arc<Mutex<ChildStdin>>,
    pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
    name:    String,
    reason:  &'static str,
    armed:   bool,
}

impl CancelOnDrop {
    fn disarm(&mut self) {
        self.armed = false;
    }
}

impl Drop for CancelOnDrop {
    fn drop(&mut self) {
        if !self.armed {
            return;
        }
        let (id, stdin, pending, name, reason) =
            (self.id, Arc::clone(&self.stdin), Arc::clone(&self.pending), self.name.clone(), self.reason);
        tokio::spawn(async move {
            pending.lock().await.remove(&id);
            debug!(target: "mcp_client", "[{name}] notifications/cancelled for request {id} ({reason})");
            write_json_line(&stdin, &crate::cancelled_notification(id, reason)).await;
        });
    }
}

/// Cooperative `tasks/cancel` for a block-and-poll `poll_task`: if the poll future
/// is dropped while still polling (a `/stop`) or hits its deadline, tell the server
/// to abandon the task. Fire-and-forget — the response carries no `pending` entry,
/// so the read-loop simply discards it. Disarmed once the task reaches a terminal
/// state (or fails), where cancelling is pointless.
struct TaskCancelOnDrop {
    request_id: u64,
    task_id:    String,
    stdin:      Arc<Mutex<ChildStdin>>,
    name:       String,
    armed:      bool,
}

impl TaskCancelOnDrop {
    fn disarm(&mut self) {
        self.armed = false;
    }
}

impl Drop for TaskCancelOnDrop {
    fn drop(&mut self) {
        if !self.armed {
            return;
        }
        let (request_id, task_id, stdin, name) =
            (self.request_id, self.task_id.clone(), Arc::clone(&self.stdin), self.name.clone());
        tokio::spawn(async move {
            debug!(target: "mcp_client", "[{name}] tasks/cancel for task {task_id}");
            write_json_line(&stdin, &crate::tasks_cancel_request(request_id, &task_id)).await;
        });
    }
}

pub struct McpServer {
    name:    String,
    stdin:   Arc<Mutex<ChildStdin>>,
    pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
    next_id: AtomicU64,
    tools:   Vec<McpTool>,
    /// Number of in-flight server→client elicitations awaiting a user reply.
    /// While > 0, an in-flight `tools/call` on this server must not time out
    /// (the user may still be typing a password into the Inbox).
    pending_elicitations: Arc<AtomicUsize>,
    /// Capabilities the server advertised in its `InitializeResult`. Captured so a
    /// future Tasks polling loop can gate on `tasks` support; unused for now.
    server_capabilities: Value,
}

impl McpServer {
    pub async fn start(
        cfg: &McpServerConfig,
        notification_tx: Option<mpsc::UnboundedSender<McpNotification>>,
        log_tx: Option<McpLogTx>,
        elicitation_handler: Option<Arc<dyn ElicitationHandler>>,
    ) -> Result<Self> {
        let command = cfg.command.as_deref()
            .ok_or_else(|| anyhow::anyhow!("stdio server '{}' requires 'command'", cfg.name))?;

        let mut cmd = Command::new(command);
        if let Some(args) = &cfg.args {
            cmd.args(args);
        }
        if let Some(env_map) = &cfg.env {
            for (k, v) in env_map {
                cmd.env(k, interpolate_env(v));
            }
        }
        cmd.stdin(Stdio::piped())
           .stdout(Stdio::piped())
           // Capture the child's stderr instead of inheriting it: many MCP
           // servers (e.g. FastMCP) print a startup banner, deprecation
           // warnings and INFO logs there, which would otherwise spill raw
           // onto our console. We drain it into tracing at `debug` level so it
           // stays quiet by default but is still available for diagnostics.
           .stderr(Stdio::piped())
           .kill_on_drop(true);

        // Detach the child into its own process group so that a terminal
        // Ctrl+C (SIGINT delivered to the whole foreground process group)
        // does not reach it directly. Otherwise Python-based MCP servers
        // catch the SIGINT and dump a KeyboardInterrupt traceback to stderr.
        // Cleanup still happens via `kill_on_drop`: when the app shuts down
        // and the reader task is dropped, the child receives SIGKILL silently.
        #[cfg(unix)]
        cmd.process_group(0);

        let mut child = cmd.spawn()
            .map_err(|e| anyhow::anyhow!("failed to spawn '{}': {e}", cfg.name))?;

        let stdin = child.stdin.take()
            .ok_or_else(|| anyhow::anyhow!("could not capture stdin for '{}'", cfg.name))?;
        // Wrap stdin early so both the struct and the read-loop (which writes
        // elicitation replies back to the server) can share the same handle.
        let stdin = Arc::new(Mutex::new(stdin));
        let stdout = child.stdout.take()
            .ok_or_else(|| anyhow::anyhow!("could not capture stdout for '{}'", cfg.name))?;

        // Drain the child's stderr into tracing at `debug` (so banners/warnings
        // don't pollute our console at the default level) *and* forward each line
        // to the host's per-server log file via `log_tx`. Per the MCP 2026 draft,
        // `stderr` is the primary place servers should log, so this is the main
        // diagnostic source for stdio transports.
        if let Some(stderr) = child.stderr.take() {
            let server_name_err = cfg.name.clone();
            let log_tx_err      = log_tx.clone();
            tokio::spawn(async move {
                let mut lines = BufReader::new(stderr).lines();
                while let Ok(Some(line)) = lines.next_line().await {
                    if !line.trim().is_empty() {
                        debug!(target: "mcp_client", "[{server_name_err}] {line}");
                        if let Some(tx) = &log_tx_err {
                            let _ = tx.send(McpLogLine::stderr(server_name_err.clone(), &line));
                        }
                    }
                }
            });
        }

        let pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>> =
            Arc::new(Mutex::new(HashMap::new()));
        let pending_elicitations = Arc::new(AtomicUsize::new(0));

        let pending_bg               = pending.clone();
        let server_name_bg           = cfg.name.clone();
        let notification_tx_bg       = notification_tx;
        let log_tx_bg                = log_tx;
        let stdin_bg                 = Arc::clone(&stdin);
        let elicitation_handler_bg   = elicitation_handler;
        let pending_elicitations_bg  = Arc::clone(&pending_elicitations);
        tokio::spawn(async move {
            let mut child = child;
            let mut lines = BufReader::new(stdout).lines();
            loop {
                match lines.next_line().await {
                    Ok(Some(line)) if !line.trim().is_empty() => {
                        if let Ok(msg) = serde_json::from_str::<Value>(&line) {
                            let has_method = msg.get("method").is_some();
                            let has_id     = msg.get("id").map(|v| !v.is_null()).unwrap_or(false);
                            if has_method && has_id {
                                // Server → client request (e.g. elicitation/create):
                                // has *both* `method` and `id`, so it must be checked
                                // before the response branch below.
                                handle_server_request(
                                    &server_name_bg, msg, &stdin_bg,
                                    &elicitation_handler_bg, &pending_elicitations_bg,
                                );
                            } else if let Some(id) = msg["id"].as_u64() {
                                if let Some(tx) = pending_bg.lock().await.remove(&id) {
                                    let _ = tx.send(msg);
                                }
                            } else if has_method {
                                // `notifications/message` is the MCP logging utility
                                // (deprecated 2026-07-28): route it to the per-server
                                // log file, not to the notification queue that feeds
                                // TIC — otherwise log records masquerade as business
                                // events. Every other notification (e.g. the custom
                                // `event/*` methods) flows on to `notification_tx`.
                                if msg.get("method").and_then(Value::as_str) == Some("notifications/message") {
                                    if let Some(tx) = &log_tx_bg {
                                        let params = &msg["params"];
                                        let _ = tx.send(McpLogLine::from_message(server_name_bg.clone(), params));
                                    }
                                } else if let Some(tx) = &notification_tx_bg {
                                    let _ = tx.send((server_name_bg.clone(), msg));
                                }
                            }
                        }
                    }
                    Ok(Some(_)) => {}
                    _ => break,
                }
            }
            let exit_info = match child.wait().await {
                Ok(status) if !status.success() => format!(
                    "process exited with {}",
                    status.code().map(|c| c.to_string()).unwrap_or_else(|| "signal".into())
                ),
                _ => "process exited unexpectedly".into(),
            };
            let error_msg = format!("MCP '{}' disconnected: {exit_info}", server_name_bg);
            if let Some(tx) = &log_tx_bg {
                let _ = tx.send(McpLogLine::lifecycle(server_name_bg.clone(), format!("disconnected: {exit_info}")));
            }
            for (_, tx) in pending_bg.lock().await.drain() {
                let _ = tx.send(json!({
                    "jsonrpc": "2.0",
                    "error": { "code": -32000, "message": error_msg }
                }));
            }
        });

        let server = McpServer {
            name:    cfg.name.clone(),
            stdin,
            pending,
            next_id: AtomicU64::new(1),
            tools:   Vec::new(),
            pending_elicitations,
            server_capabilities: json!({}),
        };

        let init = server.request("initialize", json!({
            // Declare the elicitation capability (form mode) so servers know they
            // may request input mid-call.
            "protocolVersion": crate::PROTOCOL_VERSION,
            "capabilities":    {
                "elicitation": {},
                // Experimental Tasks marker: we *recognise* a CreateTaskResult
                // (see McpCallResult::Task) but don't poll yet, so we deliberately
                // avoid claiming the full `tasks` capability — that could make a
                // server defer results we can't retrieve. Kept under `experimental`
                // until the polling loop lands.
                "experimental": { "tasks": {} },
            },
            "clientInfo":      { "name": "skald", "version": env!("CARGO_PKG_VERSION") },
        })).await?;
        // Tolerate a server that negotiates a different (older) version — warn but
        // keep going rather than disconnecting.
        if let Some(v) = init["protocolVersion"].as_str() {
            if v != crate::PROTOCOL_VERSION {
                warn!("MCP '{}': server negotiated protocol {v} (we requested {}); proceeding",
                    server.name, crate::PROTOCOL_VERSION);
            }
        }
        // Capture the server's advertised capabilities for a future Tasks poller.
        let server_capabilities = init.get("capabilities").cloned().unwrap_or_else(|| json!({}));

        server.notify("notifications/initialized", json!({})).await?;

        // Follow `nextCursor` across pages so servers with large tool lists aren't
        // silently truncated; capped at `MAX_TOOL_PAGES` against a stuck cursor.
        let mut tools: Vec<McpTool> = Vec::new();
        let mut cursor: Option<String> = None;
        for page_n in 0..crate::MAX_TOOL_PAGES {
            let params = match &cursor {
                Some(c) => json!({ "cursor": c }),
                None    => json!({}),
            };
            let page = server.request("tools/list", params).await?;
            if let Some(arr) = page["tools"].as_array() {
                tools.extend(arr.iter().map(|t| McpTool::from_json(&cfg.name, t)));
            }
            cursor = page["nextCursor"].as_str().filter(|s| !s.is_empty()).map(str::to_string);
            if cursor.is_none() {
                break;
            }
            if page_n + 1 == crate::MAX_TOOL_PAGES {
                warn!("MCP '{}': tools/list hit {}-page cap; some tools may be omitted",
                    server.name, crate::MAX_TOOL_PAGES);
            }
        }

        Ok(McpServer { tools, server_capabilities, ..server })
    }

    pub fn tools(&self) -> &[McpTool] {
        &self.tools
    }

    /// Capabilities the server advertised at `initialize`. Exposed for a future
    /// Tasks polling loop to gate on `tasks` support.
    pub fn server_capabilities(&self) -> &Value {
        &self.server_capabilities
    }

    pub async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> {
        let mut params = json!({ "name": name, "arguments": args });
        if self.wants_task(name) {
            // Opt into deferred execution for a task-capable tool (experimental
            // Tasks). Per spec, adding the `task` field to the request is the
            // opt-in for `tools/call` (the client is the requestor).
            params["task"] = json!({});
        }
        let result = self.request("tools/call", params).await?;

        if result["isError"].as_bool().unwrap_or(false) {
            anyhow::bail!("MCP tool error: {}", extract_text(&result));
        }
        match crate::extract_call_result(&result) {
            McpCallResult::Task(task) => self.poll_task(task).await,
            other => Ok(other),
        }
    }

    /// True when tool `name` advertises `execution.taskSupport` as `required`/
    /// `optional`, so we opt into deferred (Task) execution.
    fn wants_task(&self, name: &str) -> bool {
        self.tools.iter()
            .find(|t| t.name == name)
            .and_then(|t| t.task_support.as_deref())
            .is_some_and(|s| s == "required" || s == "optional")
    }

    /// Drives a deferred Task to completion (experimental Tasks, block-and-poll):
    /// polls `tasks/get` until a terminal status, then fetches the real result via
    /// `tasks/result`. A [`TaskCancelOnDrop`] guard sends `tasks/cancel` if this
    /// future is dropped (a `/stop`) or the deadline is hit. Each poll request is a
    /// normal short `request()` (subject to the 120s timeout); the *overall* wait is
    /// bounded only by the task's `ttl` — so long tasks no longer hit that wall.
    async fn poll_task(&self, task: crate::CreateTaskResult) -> Result<McpCallResult> {
        let task_id  = task.task_id.as_str();
        let cancel_id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let mut guard = TaskCancelOnDrop {
            request_id: cancel_id,
            task_id:    task.task_id.clone(),
            stdin:      Arc::clone(&self.stdin),
            name:       self.name.clone(),
            armed:      true,
        };

        let deadline     = crate::poll_deadline(task.ttl_ms);
        let mut interval = task.poll_interval_ms;

        loop {
            tokio::time::sleep(crate::clamp_poll_interval(interval)).await;
            if std::time::Instant::now() >= deadline {
                anyhow::bail!("MCP '{}' task '{task_id}' exceeded max wait", self.name);
            }
            let get = self.request("tasks/get", json!({ "taskId": task_id })).await?;
            let Some(state) = crate::CreateTaskResult::parse(&get) else {
                anyhow::bail!("MCP '{}' task '{task_id}': malformed tasks/get response", self.name);
            };
            interval = state.poll_interval_ms.or(interval);
            match state.status {
                crate::TaskStatus::Working    => continue,
                crate::TaskStatus::Completed  => break,
                crate::TaskStatus::Failed => {
                    guard.disarm();
                    anyhow::bail!("MCP '{}' task '{task_id}' failed: {}", self.name, extract_text(&get));
                }
                crate::TaskStatus::Cancelled => {
                    guard.disarm();
                    anyhow::bail!("MCP '{}' task '{task_id}' was cancelled by the server", self.name);
                }
                // Still alive but blocked on input we can't supply mid-task: abandon
                // it (guard stays armed → tasks/cancel). See follow-up.
                crate::TaskStatus::InputRequired =>
                    anyhow::bail!("MCP '{}' task '{task_id}' requires input mid-task, which isn't supported yet", self.name),
            }
        }

        // Task is terminal (completed) — nothing left to cancel.
        guard.disarm();
        let result = self.request("tasks/result", json!({ "taskId": task_id })).await?;

        if result["isError"].as_bool().unwrap_or(false) {
            anyhow::bail!("MCP tool error: {}", extract_text(&result));
        }
        Ok(crate::extract_call_result(&result))
    }

    async fn request(&self, method: &str, params: Value) -> Result<Value> {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let (tx, rx) = oneshot::channel();
        self.pending.lock().await.insert(id, tx);

        // Arm a cancellation guard for cancellable operations only (`tools/call`):
        // if this future is dropped by a `/stop` or times out before the server
        // replies, the guard notifies the server. Disarmed on a real reply.
        let mut cancel_guard = (method == "tools/call").then(|| CancelOnDrop {
            id,
            stdin:   Arc::clone(&self.stdin),
            pending: Arc::clone(&self.pending),
            name:    self.name.clone(),
            reason:  "cancelled by client",
            armed:   true,
        });

        let msg = json!({
            "jsonrpc": "2.0",
            "id":      id,
            "method":  method,
            "params":  params,
        });
        let mut line = serde_json::to_string(&msg)?;
        line.push('\n');
        self.stdin.lock().await.write_all(line.as_bytes()).await?;

        // Wait for the response, but re-arm the timeout while an elicitation is
        // in flight on this server: the user may still be typing a secret into
        // the Inbox, and the server won't answer `tools/call` until then.
        tokio::pin!(rx);
        let response = loop {
            tokio::select! {
                res = &mut rx => match res {
                    Ok(v) => break v,
                    Err(_) => {
                        // Server disconnected: nothing to cancel.
                        if let Some(g) = cancel_guard.as_mut() { g.disarm(); }
                        anyhow::bail!("MCP '{}' disconnected", self.name);
                    }
                },
                _ = tokio::time::sleep(Duration::from_secs(CALL_TIMEOUT_SECS)) => {
                    if self.pending_elicitations.load(Ordering::SeqCst) == 0 {
                        // Leave the guard armed so it fires `notifications/cancelled`;
                        // tag the reason as a timeout.
                        if let Some(g) = cancel_guard.as_mut() { g.reason = "timeout"; }
                        anyhow::bail!("MCP '{}' timed out on '{method}'", self.name);
                    }
                    // Elicitation pending: keep waiting for the user.
                }
            }
        };

        // Server replied (even with an error): the request completed — disarm.
        if let Some(g) = cancel_guard.as_mut() { g.disarm(); }

        if let Some(error) = response.get("error") {
            anyhow::bail!("MCP '{}' protocol error: {error}", self.name);
        }
        Ok(response["result"].clone())
    }

    async fn notify(&self, method: &str, params: Value) -> Result<()> {
        let msg = json!({
            "jsonrpc": "2.0",
            "method":  method,
            "params":  params,
        });
        let mut line = serde_json::to_string(&msg)?;
        line.push('\n');
        self.stdin.lock().await.write_all(line.as_bytes()).await?;
        Ok(())
    }
}

#[async_trait]
impl McpServerClient for McpServer {
    fn tools(&self) -> &[McpTool] { self.tools() }
    async fn call_tool(&self, name: &str, args: Value) -> Result<McpCallResult> { self.call_tool(name, args).await }
}