//! JSON-RPC 2.0 over an ACP agent's stdio: newline-delimited messages, the
//! agent's requests answered here, its notifications queued for the turn.

use std::collections::HashMap;
use std::path::Path;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, PoisonError};

use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, Command};
use tokio::sync::{mpsc, oneshot};

use super::provider_key::KeyVar;
use super::stream::mcp_server_of;
use super::{AcpError, PermissionAsk, PermissionPolicy};
use crate::spec::AcpCommand;

/// Something the agent sent that is not a response to one of our requests.
pub(super) enum Incoming {
    /// The `params` of a `session/update` notification.
    Update(Value),
    Stdout(String),
    Stderr(String),
    /// The agent closed its stdout: it has exited or is about to.
    Closed,
}

type Reply = Result<Value, AcpError>;
type Pending = Arc<Mutex<HashMap<u64, (&'static str, oneshot::Sender<Reply>)>>>;

pub(super) struct Connection {
    child: Child,
    stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
    next_id: AtomicU64,
    pending: Pending,
    pub(super) incoming: mpsc::UnboundedReceiver<Incoming>,
}

impl Connection {
    /// Spawn the agent in `cwd`, with `key` added to its environment, and
    /// start reading its output.
    pub(super) fn spawn(
        command: &AcpCommand,
        key: Option<&KeyVar>,
        cwd: &Path,
        permit: PermissionPolicy,
        servers: Vec<String>,
    ) -> Result<Self, AcpError> {
        let mut cmd = Command::new(&command.command);
        cmd.args(&command.args).envs(&command.env);
        if let Some(key) = key {
            cmd.env(&key.name, &key.value);
        }
        let mut child = cmd
            .current_dir(cwd)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true)
            .spawn()
            .map_err(|source| AcpError::Spawn {
                program: command.command.clone(),
                source,
            })?;
        let (Some(stdin), Some(stdout), Some(stderr)) =
            (child.stdin.take(), child.stdout.take(), child.stderr.take())
        else {
            unreachable!("all three stdio handles were requested as pipes");
        };
        let stdin = Arc::new(tokio::sync::Mutex::new(stdin));
        let pending: Pending = Arc::default();
        let (tx, incoming) = mpsc::unbounded_channel();

        let stderr_tx = tx.clone();
        tokio::spawn(async move {
            let mut lines = BufReader::new(stderr).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                if stderr_tx.send(Incoming::Stderr(line)).is_err() {
                    break;
                }
            }
        });

        let reader = Reader {
            stdin: stdin.clone(),
            pending: pending.clone(),
            tx,
            permit,
            servers,
        };
        tokio::spawn(async move {
            let mut lines = BufReader::new(stdout).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                reader.dispatch(&line).await;
            }
            // Fail every request still waiting: no reply can come now.
            reader
                .pending
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .clear();
            let _ = reader.tx.send(Incoming::Closed);
        });

        Ok(Self {
            child,
            stdin,
            next_id: AtomicU64::new(1),
            pending,
            incoming,
        })
    }

    /// Whether the agent process has exited.
    pub(super) fn exited(&mut self) -> bool {
        !matches!(self.child.try_wait(), Ok(None))
    }

    /// Send a request and return the receiver its reply arrives on.
    pub(super) async fn send(
        &self,
        method: &'static str,
        params: Value,
    ) -> Result<oneshot::Receiver<Reply>, AcpError> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = oneshot::channel();
        self.pending
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .insert(id, (method, tx));
        let message = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
        if let Err(e) = write(&self.stdin, &message).await {
            self.pending
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .remove(&id);
            return Err(e);
        }
        Ok(rx)
    }

    /// Send a notification: a message the agent does not reply to.
    pub(super) async fn notify(&self, method: &'static str, params: Value) -> Result<(), AcpError> {
        let message = json!({ "jsonrpc": "2.0", "method": method, "params": params });
        write(&self.stdin, &message).await
    }

    /// Send a request and wait for its reply, queueing anything else the agent
    /// sends meanwhile on [`Self::incoming`].
    pub(super) async fn request(
        &self,
        method: &'static str,
        params: Value,
    ) -> Result<Value, AcpError> {
        let rx = self.send(method, params).await?;
        rx.await.unwrap_or(Err(AcpError::Closed))
    }
}

/// The stdout side: routes responses to their waiting request, answers the
/// agent's own requests, and queues its notifications.
struct Reader {
    stdin: Arc<tokio::sync::Mutex<ChildStdin>>,
    pending: Pending,
    tx: mpsc::UnboundedSender<Incoming>,
    permit: PermissionPolicy,
    servers: Vec<String>,
}

impl Reader {
    async fn dispatch(&self, line: &str) {
        let Ok(message) = serde_json::from_str::<Value>(line) else {
            if !line.trim().is_empty() {
                let _ = self.tx.send(Incoming::Stdout(line.to_owned()));
            }
            return;
        };
        let method = message.get("method").and_then(Value::as_str);
        let id = message.get("id");
        match (method, id) {
            (None, Some(id)) => self.resolve(id, &message),
            (Some(method), Some(id)) => {
                let reply = match method {
                    "session/request_permission" => json!({
                        "jsonrpc": "2.0", "id": id,
                        "result": permission_outcome(&message["params"], &self.permit, &self.servers),
                    }),
                    _ => json!({
                        "jsonrpc": "2.0", "id": id,
                        "error": { "code": -32601, "message": format!("method not found: {method}") },
                    }),
                };
                if let Err(e) = write(&self.stdin, &reply).await {
                    tracing::warn!(error = %e, method, "failed to answer ACP agent request");
                }
            }
            (Some("session/update"), None) => {
                let params = message.get("params").cloned().unwrap_or(Value::Null);
                let _ = self.tx.send(Incoming::Update(params));
            }
            _ => {}
        }
    }

    fn resolve(&self, id: &Value, message: &Value) {
        let Some(id) = id.as_u64() else { return };
        let Some((method, tx)) = self
            .pending
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .remove(&id)
        else {
            return;
        };
        let reply = match message.get("error") {
            Some(error) => Err(AcpError::Rpc {
                method,
                code: error.get("code").and_then(Value::as_i64).unwrap_or(0),
                message: error
                    .get("message")
                    .and_then(Value::as_str)
                    .unwrap_or("")
                    .to_owned(),
                data: error.get("data").map(Value::to_string).unwrap_or_default(),
            }),
            None => Ok(message.get("result").cloned().unwrap_or(Value::Null)),
        };
        let _ = tx.send(reply);
    }
}

async fn write(stdin: &tokio::sync::Mutex<ChildStdin>, message: &Value) -> Result<(), AcpError> {
    let mut line = message.to_string();
    line.push('\n');
    let mut stdin = stdin.lock().await;
    stdin.write_all(line.as_bytes()).await?;
    stdin.flush().await?;
    Ok(())
}

/// The `session/request_permission` result: the agent's first option of the
/// matching kind, allow if `permit` accepts the request and reject otherwise.
/// With no such option the request is answered `cancelled`. `servers` are the
/// MCP servers handed to the session, matched against the tool call's title.
pub(super) fn permission_outcome(
    params: &Value,
    permit: &PermissionPolicy,
    servers: &[String],
) -> Value {
    let tool_call = params.get("toolCall");
    let field = |k: &str| tool_call.and_then(|t| t.get(k)).and_then(Value::as_str);
    let kind = field("kind").unwrap_or("other");
    let ask = PermissionAsk {
        kind,
        // MCP tool calls are kind `other`; a title alone must not lift another
        // kind (`execute`, `fetch`) into one.
        mcp_server: (kind == "other")
            .then(|| field("title").and_then(|title| mcp_server_of(title, servers)))
            .flatten(),
    };
    let wanted = if permit(&ask) {
        ["allow_once", "allow_always"]
    } else {
        ["reject_once", "reject_always"]
    };
    let options = params
        .get("options")
        .and_then(Value::as_array)
        .map_or(&[][..], Vec::as_slice);
    let chosen = wanted.iter().find_map(|want| {
        options
            .iter()
            .find(|o| o.get("kind").and_then(Value::as_str) == Some(want))
            .and_then(|o| o.get("optionId"))
    });
    match chosen {
        Some(option) => json!({ "outcome": { "outcome": "selected", "optionId": option } }),
        None => json!({ "outcome": { "outcome": "cancelled" } }),
    }
}

#[cfg(test)]
mod tests {
    use super::permission_outcome;
    use crate::acp::{PermissionAsk, PermissionPolicy};
    use serde_json::json;
    use std::sync::Arc;

    fn request(kind: &str, title: &str) -> serde_json::Value {
        json!({
            "sessionId": "s",
            "toolCall": { "toolCallId": "t", "kind": kind, "title": title },
            "options": [
                { "optionId": "once", "kind": "allow_once", "name": "Allow once" },
                { "optionId": "always", "kind": "allow_always", "name": "Always allow" },
                { "optionId": "reject", "kind": "reject_once", "name": "Reject" },
            ],
        })
    }

    fn servers() -> Vec<String> {
        vec!["hyperhive".to_owned()]
    }

    /// Allows `edit` and MCP tools, nothing else.
    fn edit_or_mcp() -> PermissionPolicy {
        Arc::new(|ask: &PermissionAsk<'_>| ask.kind == "edit" || ask.mcp_server.is_some())
    }

    #[test]
    fn a_permitted_request_is_allowed_once() {
        assert_eq!(
            permission_outcome(&request("edit", "edit"), &edit_or_mcp(), &servers()),
            json!({ "outcome": { "outcome": "selected", "optionId": "once" } })
        );
    }

    #[test]
    fn a_refused_request_is_rejected() {
        assert_eq!(
            permission_outcome(&request("execute", "bash"), &edit_or_mcp(), &servers()),
            json!({ "outcome": { "outcome": "selected", "optionId": "reject" } })
        );
    }

    #[test]
    fn the_policy_sees_which_mcp_server_a_tool_belongs_to() {
        let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
        let record = seen.clone();
        let policy: PermissionPolicy = Arc::new(move |ask: &PermissionAsk<'_>| {
            record
                .lock()
                .unwrap()
                .push((ask.kind.to_owned(), ask.mcp_server.map(str::to_owned)));
            true
        });
        permission_outcome(&request("other", "hyperhive_send"), &policy, &servers());
        permission_outcome(&request("other", "skill"), &policy, &servers());
        let mut req = request("x", "read");
        req["toolCall"].as_object_mut().unwrap().remove("kind");
        permission_outcome(&req, &policy, &servers());
        assert_eq!(
            *seen.lock().unwrap(),
            vec![
                ("other".to_owned(), Some("hyperhive".to_owned())),
                ("other".to_owned(), None),
                ("other".to_owned(), None),
            ]
        );
    }

    /// A title shaped like `<server>_<tool>` marks an MCP call only on kind
    /// `other`, which is what an agent gives an MCP tool call.
    #[test]
    fn an_mcp_looking_title_on_execute_or_fetch_is_not_an_mcp_call() {
        let policy: PermissionPolicy = Arc::new(|ask: &PermissionAsk<'_>| ask.mcp_server.is_some());
        for kind in ["execute", "fetch"] {
            assert_eq!(
                permission_outcome(&request(kind, "hyperhive_send"), &policy, &servers())["outcome"]
                    ["optionId"],
                "reject",
                "{kind}"
            );
        }
        assert_eq!(
            permission_outcome(&request("other", "hyperhive_send"), &policy, &servers())["outcome"]
                ["optionId"],
            "once"
        );
    }

    #[test]
    fn no_matching_option_cancels() {
        let req = json!({ "toolCall": { "kind": "execute" },
                          "options": [{ "optionId": "once", "kind": "allow_once" }] });
        assert_eq!(
            permission_outcome(&req, &edit_or_mcp(), &servers()),
            json!({ "outcome": { "outcome": "cancelled" } })
        );
    }
}