//! The ACP backend: one long-lived agent process per runtime, spawned on the
//! first turn, holding one durable session whose id is kept in a file.

mod rpc;
mod stream;

use std::collections::VecDeque;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use hive_claude::{Config, Progress, Sink};
use serde_json::{Value, json};
use tokio::sync::Mutex;

use self::rpc::{Connection, Incoming};
use self::stream::StreamMapper;
use crate::spec::AcpCommand;
use crate::{Error, Result, Runtime};

/// ACP protocol version this client speaks.
const PROTOCOL_VERSION: u64 = 1;

/// Stderr lines kept for the error when the agent exits mid-turn.
const STDERR_TAIL: usize = 20;

/// After a turn's reply, how long the agent must stay quiet before the turn
/// is taken to be over, and the most that wait may take in all.
const SETTLE_QUIET: std::time::Duration = std::time::Duration::from_millis(500);
const SETTLE_MAX: std::time::Duration = std::time::Duration::from_secs(5);

/// One `session/request_permission` request, as a [`PermissionPolicy`] sees it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PermissionAsk<'a> {
    /// The tool call's ACP `kind` (`read`, `edit`, `execute`, `fetch`, …;
    /// `other` when the agent gives none).
    pub kind: &'a str,
    /// For a `kind` of `other`: the MCP server, of those handed to the
    /// session, whose tool this is — the tool call's title names it as
    /// `<server>_<tool>`, `<server>__<tool>` or `mcp__<server>__<tool>`.
    /// `None` for every other kind.
    pub mcp_server: Option<&'a str>,
}

/// Decides the agent's `session/request_permission` requests. `true` allows
/// the call once, `false` rejects it.
pub type PermissionPolicy = Arc<dyn Fn(&PermissionAsk<'_>) -> bool + Send + Sync>;

/// Why an ACP operation failed.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum AcpError {
    #[error("failed to spawn ACP agent `{program}`: {source}")]
    Spawn {
        program: String,
        #[source]
        source: std::io::Error,
    },
    #[error("ACP agent closed its output (exited?)")]
    Closed,
    #[error("ACP agent exited mid-turn; last stderr:\n{stderr_tail}")]
    Exited { stderr_tail: String },
    #[error("ACP agent i/o failed: {0}")]
    Io(#[from] std::io::Error),
    #[error("ACP `{method}` failed: {message} (code {code}) {data}")]
    Rpc {
        method: &'static str,
        code: i64,
        message: String,
        data: String,
    },
    #[error("ACP `{method}` response has no {field}")]
    Malformed {
        method: &'static str,
        field: &'static str,
    },
    #[error(
        "the ACP agent does not accept MCP servers over HTTP (`mcpCapabilities.http` is not true), \
         and the hyperhive tools are only served over HTTP"
    )]
    NoHttpMcp,
    #[error("reading the MCP config {path}: {detail}")]
    McpConfig { path: PathBuf, detail: String },
}

/// Turns on an ACP agent's durable session.
///
/// From the [`Config`] it reads `cwd` (the agent's working directory),
/// `mcp_config` (converted to the session's `mcpServers`) and
/// `system_prompt_file` (prepended to the first prompt of each new session,
/// since ACP has no system prompt). Everything else in it is claude's and is
/// ignored, including `idle_timeout`.
pub struct AcpRuntime {
    command: AcpCommand,
    session_file: PathBuf,
    permit: PermissionPolicy,
    live: Mutex<Option<Live>>,
}

/// The running agent process and the session loaded into it.
struct Live {
    conn: Connection,
    load_session: bool,
    /// The recorded session, or a new one whose first prompt is unanswered.
    loaded: Option<String>,
    model: Option<String>,
}

impl AcpRuntime {
    /// A runtime spawning `command`, keeping its session id in
    /// `session_file`, and answering its permission requests with `permit`.
    #[must_use]
    pub fn new(command: AcpCommand, session_file: PathBuf, permit: PermissionPolicy) -> Self {
        Self {
            command,
            session_file,
            permit,
            live: Mutex::new(None),
        }
    }

    /// Names a new session until its first prompt has been answered, so a
    /// failed first prompt is retried on the same session.
    fn pending_file(&self) -> PathBuf {
        let mut path = self.session_file.clone().into_os_string();
        path.push(".pending");
        PathBuf::from(path)
    }

    async fn start(&self, config: &Config) -> Result<Live> {
        let (_, servers) = mcp_servers(config)?;
        let cwd = session_cwd(config);
        let conn = Connection::spawn(&self.command, &cwd, self.permit.clone(), servers)?;
        let init = conn
            .request(
                "initialize",
                json!({
                    "protocolVersion": PROTOCOL_VERSION,
                    "clientCapabilities": {
                        "fs": { "readTextFile": false, "writeTextFile": false },
                        "terminal": false,
                    },
                    "clientInfo": { "name": "hyperhive", "version": env!("CARGO_PKG_VERSION") },
                }),
            )
            .await?;
        let caps = &init["agentCapabilities"];
        if caps["mcpCapabilities"]["http"] != Value::Bool(true) {
            return Err(AcpError::NoHttpMcp.into());
        }
        tracing::info!(agent = %init["agentInfo"], "ACP agent initialized");
        Ok(Live {
            conn,
            load_session: caps["loadSession"] == Value::Bool(true),
            loaded: None,
            model: None,
        })
    }

    /// Attach the durable session: the recorded one (named by the session
    /// file), else a new one whose first prompt was never answered (named by
    /// the pending file), else a new one. Returns its id and whether its first
    /// prompt, the one carrying the system prompt, is still to be answered;
    /// [`Self::turn`] records the session once it has been.
    async fn attach(
        &self,
        live: &mut Live,
        cwd: &Path,
        servers: &[Value],
    ) -> Result<(String, bool)> {
        let known = [
            (read_id(&self.session_file), false),
            (read_id(&self.pending_file()), true),
        ]
        .into_iter()
        .filter_map(|(id, new)| Some((id?, new)));
        for (id, new) in known {
            if live.loaded.as_deref() == Some(id.as_str()) {
                return Ok((id, new));
            }
            if !live.load_session {
                continue;
            }
            let params = json!({ "sessionId": id, "cwd": cwd, "mcpServers": servers });
            match live.conn.request("session/load", params).await {
                Ok(response) => {
                    discard_stale(&mut live.conn)?;
                    live.model = stream::session_model(&response);
                    live.loaded = Some(id.clone());
                    return Ok((id, new));
                }
                Err(e @ AcpError::Rpc { .. }) => {
                    tracing::warn!(error = %e, session = %id, "ACP session/load failed");
                }
                Err(e) => return Err(e.into()),
            }
        }
        let params = json!({ "cwd": cwd, "mcpServers": servers });
        let response = live.conn.request("session/new", params).await?;
        let id = response["sessionId"]
            .as_str()
            .ok_or(AcpError::Malformed {
                method: "session/new",
                field: "sessionId",
            })?
            .to_owned();
        live.model = stream::session_model(&response);
        live.loaded = Some(id.clone());
        write_id(&self.pending_file(), &id)?;
        Ok((id, true))
    }

    fn record_session(&self, id: &str) -> std::result::Result<(), AcpError> {
        write_id(&self.session_file, id)?;
        remove_if_present(&self.pending_file())
    }

    async fn turn(
        &self,
        live: &mut Live,
        config: &Config,
        prompt: &str,
        sink: &impl Sink,
    ) -> Result<Progress> {
        let cwd = session_cwd(config);
        let (servers, names) = mcp_servers(config)?;
        let (session, created) = self.attach(live, &cwd, &servers).await?;
        let text = match (created, &config.system_prompt_file) {
            (true, Some(path)) => match std::fs::read_to_string(path) {
                Ok(system) => format!("{system}\n\n{prompt}"),
                Err(e) => {
                    tracing::warn!(path = %path.display(), error = %e, "system prompt unreadable");
                    prompt.to_owned()
                }
            },
            _ => prompt.to_owned(),
        };
        let params = json!({ "sessionId": session, "prompt": [{ "type": "text", "text": text }] });
        discard_stale(&mut live.conn)?;
        let mut reply = live.conn.send("session/prompt", params).await?;
        let mut mapper = StreamMapper::new(names);
        let mut stderr_tail = VecDeque::new();
        let response = loop {
            tokio::select! {
                // Notifications first: the ones sent before the reply belong
                // to this turn and must reach the sink before it ends.
                biased;
                incoming = live.conn.incoming.recv() => {
                    let incoming = incoming.unwrap_or(Incoming::Closed);
                    if !deliver(incoming, &session, &mut mapper, &mut stderr_tail, sink) {
                        let stderr_tail = Vec::from(stderr_tail).join("\n");
                        return Err(AcpError::Exited { stderr_tail }.into());
                    }
                }
                reply = &mut reply => break reply.unwrap_or(Err(AcpError::Closed))?,
            }
        };
        if created {
            self.record_session(&session)?;
        }
        // Updates may still arrive after the reply: an agent can forward its
        // stream asynchronously. Take them until the agent goes quiet.
        let settled = tokio::time::Instant::now() + SETTLE_MAX;
        while let Ok(Some(incoming)) = tokio::time::timeout_at(
            settled.min(tokio::time::Instant::now() + SETTLE_QUIET),
            live.conn.incoming.recv(),
        )
        .await
        {
            if !deliver(incoming, &session, &mut mapper, &mut stderr_tail, sink) {
                break;
            }
        }
        for event in mapper.finish() {
            sink.on_event(&event);
        }
        match response["stopReason"].as_str() {
            Some("end_turn") => {}
            reason => {
                let reason = reason.unwrap_or("none");
                tracing::warn!(reason, "ACP turn stopped early");
                sink.on_stderr_line(&format!("ACP turn stopped: {reason}"));
            }
        }
        let mut progress = Progress::default();
        progress.created = created;
        progress.telemetry = mapper.telemetry(&response, live.model.as_deref());
        Ok(progress)
    }
}

impl Runtime for AcpRuntime {
    async fn run(&self, config: &Config, prompt: &str, sink: &impl Sink) -> Result<Progress> {
        let mut guard = self.live.lock().await;
        if let Some(live) = guard.as_mut()
            && live.conn.exited()
        {
            tracing::warn!("ACP agent has exited; respawning");
            *guard = None;
        }
        let live = match guard.as_mut() {
            Some(live) => live,
            None => guard.insert(self.start(config).await?),
        };
        let result = self.turn(live, config, prompt, sink).await;
        if let Err(Error::Acp(
            AcpError::Closed
            | AcpError::Exited { .. }
            | AcpError::Io(_)
            | AcpError::Malformed { .. },
        )) = &result
        {
            *guard = None;
        }
        result
    }

    async fn compact(&self, _config: &Config, _sink: &impl Sink) -> Result<()> {
        Err(Error::Unsupported("compact"))
    }

    /// Moves the session file aside; the agent keeps its own copy of the
    /// session, so nothing is lost. A new session whose first prompt was never
    /// answered is dropped too.
    fn archive(&self) -> Result<Option<PathBuf>> {
        remove_if_present(&self.pending_file()).map_err(Error::Acp)?;
        if !self.session_file.exists() {
            return Ok(None);
        }
        let stamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs());
        let mut archived = self.session_file.clone().into_os_string();
        archived.push(format!(".{stamp}.archived"));
        let archived = PathBuf::from(archived);
        std::fs::rename(&self.session_file, &archived).map_err(AcpError::Io)?;
        Ok(Some(archived))
    }
}

/// Hand one incoming item to the sink. Returns `false` once the agent has
/// closed its output.
fn deliver(
    incoming: Incoming,
    session: &str,
    mapper: &mut StreamMapper,
    stderr_tail: &mut VecDeque<String>,
    sink: &impl Sink,
) -> bool {
    match incoming {
        Incoming::Update(params) => {
            if params["sessionId"].as_str() == Some(session) {
                for event in mapper.push(&params["update"]) {
                    sink.on_event(&event);
                }
            }
        }
        Incoming::Stdout(line) => sink.on_stdout_line(&line),
        Incoming::Stderr(line) => {
            sink.on_stderr_line(&line);
            if stderr_tail.len() == STDERR_TAIL {
                stderr_tail.pop_front();
            }
            stderr_tail.push_back(line);
        }
        Incoming::Closed => return false,
    }
    true
}

/// Drop what the agent sent outside a turn: the history a `session/load`
/// replays (the caller already has it), and anything that arrived after the
/// previous turn settled, which must not be shown as part of the next one.
fn discard_stale(conn: &mut Connection) -> std::result::Result<(), AcpError> {
    while let Ok(incoming) = conn.incoming.try_recv() {
        match incoming {
            Incoming::Update(_) => {}
            Incoming::Stdout(line) | Incoming::Stderr(line) => {
                tracing::info!(line = %line, "ACP agent output between turns");
            }
            Incoming::Closed => return Err(AcpError::Closed),
        }
    }
    Ok(())
}

fn read_id(file: &Path) -> Option<String> {
    std::fs::read_to_string(file)
        .ok()
        .map(|s| s.trim().to_owned())
        .filter(|s| !s.is_empty())
}

fn write_id(file: &Path, id: &str) -> std::result::Result<(), AcpError> {
    if let Some(dir) = file.parent() {
        std::fs::create_dir_all(dir)?;
    }
    std::fs::write(file, id)?;
    Ok(())
}

fn remove_if_present(file: &Path) -> std::result::Result<(), AcpError> {
    match std::fs::remove_file(file) {
        Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e.into()),
        _ => Ok(()),
    }
}

fn session_cwd(config: &Config) -> PathBuf {
    config
        .cwd
        .clone()
        .or_else(|| std::env::current_dir().ok())
        .unwrap_or_else(|| PathBuf::from("."))
}

fn mcp_servers(config: &Config) -> Result<(Vec<Value>, Vec<String>)> {
    let Some(path) = &config.mcp_config else {
        return Ok((Vec::new(), Vec::new()));
    };
    let fail = |detail: String| AcpError::McpConfig {
        path: path.clone(),
        detail,
    };
    let raw = std::fs::read_to_string(path).map_err(|e| fail(e.to_string()))?;
    let parsed: Value = serde_json::from_str(&raw).map_err(|e| fail(e.to_string()))?;
    Ok(stream::mcp_servers(&parsed))
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;
    use std::path::Path;
    use std::sync::Arc;

    use hive_claude::{Config, NoopSink};

    use super::{AcpError, AcpRuntime};
    use crate::{AcpCommand, Error, Runtime};

    /// An ACP agent in plain `sh`. It answers `initialize`, numbers its
    /// sessions `s1`, `s2`, …, appends every method it is sent to
    /// `<log>.methods` (and `start` when it starts), and every
    /// `session/prompt` line to `<log>`. It answers every prompt with
    /// `end_turn`, except the very first one it is sent across restarts,
    /// which it handles per its mode:
    ///
    /// - `fail`: an error.
    const AGENT: &str = r#"
n=0
printf 'start\n' >> "$1.methods"
while IFS= read -r line; do
  id=${line#*\"id\":}; id=${id%%[,\}]*}
  m=${line#*\"method\":\"}; m=${m%%\"*}
  printf '%s\n' "$m" >> "$1.methods"
  case $m in
    initialize)
      printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true,"mcpCapabilities":{"http":true}}}}\n' "$id" ;;
    session/new)
      n=$((n+1))
      printf '{"jsonrpc":"2.0","id":%s,"result":{"sessionId":"s%s"}}\n' "$id" "$n" ;;
    session/load)
      printf '{"jsonrpc":"2.0","id":%s,"result":{}}\n' "$id" ;;
    session/prompt)
      printf '%s\n' "$line" >> "$1"
      if [ -e "$1.once" ]; then
        printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$id"
      else
        : > "$1.once"
        case $2 in
          fail)
            printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32603,"message":"provider down"}}\n' "$id" ;;
        esac
      fi ;;
  esac
done
"#;

    fn runtime(dir: &Path, mode: &str) -> AcpRuntime {
        let script = dir.join("agent.sh");
        std::fs::write(&script, AGENT).unwrap();
        let command = AcpCommand {
            command: "/bin/sh".into(),
            args: vec![
                script.display().to_string(),
                dir.join("prompts").display().to_string(),
                mode.to_owned(),
            ],
            env: BTreeMap::new(),
        };
        AcpRuntime::new(command, dir.join("session"), Arc::new(|_: &_| false))
    }

    fn config(dir: &Path) -> Config {
        std::fs::write(dir.join("system.md"), "SYSTEM PROMPT").unwrap();
        Config {
            cwd: Some(dir.to_path_buf()),
            system_prompt_file: Some(dir.join("system.md")),
            ..Config::default()
        }
    }

    /// Which of the prompts the agent was sent carried the system prompt.
    fn carried_system_prompt(dir: &Path) -> Vec<bool> {
        std::fs::read_to_string(dir.join("prompts"))
            .unwrap()
            .lines()
            .map(|line| line.contains("SYSTEM PROMPT"))
            .collect()
    }

    /// How often the agent started, or was sent `method`.
    fn count(dir: &Path, method: &str) -> usize {
        std::fs::read_to_string(dir.join("prompts.methods"))
            .unwrap()
            .lines()
            .filter(|line| *line == method)
            .count()
    }

    fn recorded(dir: &Path) -> String {
        std::fs::read_to_string(dir.join("session")).unwrap()
    }

    #[tokio::test]
    async fn a_session_whose_first_prompt_failed_is_retried_not_resumed() {
        let dir = tempfile::tempdir().unwrap();
        let (runtime, config) = (runtime(dir.path(), "fail"), config(dir.path()));

        let failed = runtime.run(&config, "one", &NoopSink).await;
        assert!(
            matches!(failed, Err(Error::Acp(AcpError::Rpc { .. }))),
            "{failed:?}"
        );
        assert!(!dir.path().join("session").exists());

        let retried = runtime.run(&config, "two", &NoopSink).await.unwrap();
        assert!(retried.created);
        let next = runtime.run(&config, "three", &NoopSink).await.unwrap();
        assert!(!next.created);

        assert_eq!(carried_system_prompt(dir.path()), [true, true, false]);
        assert_eq!(recorded(dir.path()), "s1");
        assert_eq!(count(dir.path(), "session/new"), 1);
        assert!(!dir.path().join("session.pending").exists());
    }

    #[tokio::test]
    async fn after_a_restart_a_session_whose_first_prompt_failed_is_retried() {
        let dir = tempfile::tempdir().unwrap();
        let config = config(dir.path());

        let before = runtime(dir.path(), "fail");
        assert!(before.run(&config, "one", &NoopSink).await.is_err());
        drop(before);

        let after = runtime(dir.path(), "fail");
        let retried = after.run(&config, "two", &NoopSink).await.unwrap();
        assert!(retried.created);

        assert_eq!(carried_system_prompt(dir.path()), [true, true]);
        assert_eq!(count(dir.path(), "start"), 2);
        assert_eq!(count(dir.path(), "session/new"), 1);
        assert_eq!(count(dir.path(), "session/load"), 1);
        assert_eq!(recorded(dir.path()), "s1");
    }
}