use anyhow::Result;
use async_stream::try_stream;
use async_trait::async_trait;
use futures::future::BoxFuture;
use rmcp::model::{Role, Tool};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use tempfile::NamedTempFile;
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::process::Command;
use tokio::sync::oneshot;

use super::base::{
    stream_from_single_message, ConfigKey, MessageStream, PermissionRouting, Provider, ProviderDef,
    ProviderMetadata, ProviderUsage, Usage,
};
use super::errors::ProviderError;
use super::utils::filter_extensions_from_system_prompt;
use crate::config::base::ClaudeCodeCommand;
use crate::config::paths::Paths;
use crate::config::search_path::SearchPaths;
use crate::config::{Config, ExtensionConfig, GooseMode};
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
use crate::permission::permission_confirmation::PrincipalType;
use crate::permission::{Permission, PermissionConfirmation};
use crate::subprocess::configure_subprocess;

use super::cli_common::{error_from_event, extract_usage_tokens};

const CLAUDE_CODE_PROVIDER_NAME: &str = "claude-code";
pub const CLAUDE_CODE_DEFAULT_MODEL: &str = "default";
pub const CLAUDE_CODE_DOC_URL: &str = "https://code.claude.com/docs/en/setup";

// https://github.com/anthropics/claude-agent-sdk-python/blob/0e9397e/src/claude_agent_sdk/types.py#L857-L859
#[derive(Serialize)]
struct ControlResponse<T: Serialize> {
    #[serde(rename = "type")]
    msg_type: &'static str,
    response: ControlResponseBody<T>,
}

#[derive(Serialize)]
struct ControlResponseBody<T: Serialize> {
    subtype: &'static str,
    request_id: String,
    response: T,
}

// https://github.com/anthropics/claude-agent-sdk-python/blob/0e9397e/src/claude_agent_sdk/types.py#L135-L153
#[derive(Serialize)]
#[serde(tag = "behavior")]
enum PermissionResponse {
    #[serde(rename = "allow")]
    Allow {
        #[serde(rename = "updatedInput")]
        updated_input: serde_json::Map<String, Value>,
        #[serde(rename = "toolUseID")]
        tool_use_id: String,
    },
    #[serde(rename = "deny")]
    Deny { message: String },
}

#[derive(Serialize)]
struct ControlRequest {
    #[serde(rename = "type")]
    msg_type: &'static str,
    request_id: String,
    request: ControlRequestBody,
}

#[derive(Serialize)]
#[serde(tag = "subtype")]
enum ControlRequestBody {
    #[serde(rename = "initialize")]
    Initialize,
    #[serde(rename = "set_model")]
    SetModel { model: String },
}

impl ControlRequestBody {
    fn label(&self) -> &'static str {
        match self {
            Self::Initialize => "initialize",
            Self::SetModel { .. } => "set_model",
        }
    }
}

#[derive(Deserialize)]
struct IncomingControlResponse {
    response: IncomingControlResponseBody,
}

#[derive(Deserialize)]
#[serde(tag = "subtype")]
enum IncomingControlResponseBody {
    #[serde(rename = "success")]
    Success {
        request_id: String,
        #[serde(default)]
        response: Option<Value>,
    },
    #[serde(rename = "error")]
    Error {
        request_id: String,
        #[serde(default)]
        error: String,
    },
}

#[derive(Deserialize)]
struct IncomingControlRequest {
    request_id: String,
    request: IncomingRequestBody,
}

#[derive(Deserialize)]
#[serde(tag = "subtype")]
enum IncomingRequestBody {
    #[serde(rename = "can_use_tool")]
    CanUseTool {
        tool_name: String,
        #[serde(default)]
        input: serde_json::Map<String, Value>,
        #[serde(default)]
        tool_use_id: String,
    },
}

impl<T: Serialize> ControlResponse<T> {
    fn success(request_id: String, response: T) -> Self {
        Self {
            msg_type: "control_response",
            response: ControlResponseBody {
                subtype: "success",
                request_id,
                response,
            },
        }
    }
}

struct CliProcess {
    child: tokio::process::Child,
    stdin: Box<dyn tokio::io::AsyncWrite + Unpin + Send>,
    reader: BufReader<Box<dyn tokio::io::AsyncRead + Unpin + Send>>,
    #[allow(dead_code)]
    stderr_handle: tokio::task::JoinHandle<String>,
    current_model: String,
    log_model_update: bool,
    next_request_id: u64,
    needs_drain: bool,
}

impl std::fmt::Debug for CliProcess {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CliProcess")
            .field("current_model", &self.current_model)
            .field("next_request_id", &self.next_request_id)
            .finish_non_exhaustive()
    }
}

impl CliProcess {
    fn next_request_id(&mut self) -> String {
        let id = self.next_request_id;
        self.next_request_id += 1;
        format!("req_{id}")
    }

    async fn send_control_request(
        &mut self,
        body: ControlRequestBody,
    ) -> Result<Option<Value>, ProviderError> {
        let request_id = self.next_request_id();
        exchange_control(&mut self.stdin, &mut self.reader, &request_id, body).await
    }

    async fn send_set_model(&mut self, model: &str) -> Result<(), ProviderError> {
        if model == self.current_model {
            return Ok(());
        }
        self.send_control_request(ControlRequestBody::SetModel {
            model: model.to_string(),
        })
        .await?;
        self.current_model = model.to_string();
        self.log_model_update = true;
        Ok(())
    }

    async fn drain_pending_response(&mut self) {
        if !self.needs_drain {
            return;
        }
        tracing::debug!("Draining cancelled response from CLI process");

        let drain = async {
            let mut line = String::new();
            loop {
                line.clear();
                match self.reader.read_line(&mut line).await {
                    Ok(0) => break,
                    Ok(_) => {
                        let trimmed = line.trim();
                        if trimmed.is_empty() {
                            continue;
                        }
                        if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
                            match parsed.get("type").and_then(|t| t.as_str()) {
                                Some("result") | Some("error") => break,
                                _ => continue,
                            }
                        } else {
                            tracing::trace!(line = trimmed, "Non-JSON line during drain");
                        }
                    }
                    Err(_) => break,
                }
            }
        };

        const DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
        if tokio::time::timeout(DRAIN_TIMEOUT, drain).await.is_err() {
            // CLI is still producing the old response. Leave needs_drain
            // true so the next call retries — by then the old response
            // likely completed and drain will succeed quickly.
            tracing::warn!(
                "Drain did not complete in {DRAIN_TIMEOUT:?}; \
                 will retry on next request"
            );
            return;
        }

        self.needs_drain = false;
        tracing::debug!("Drain complete, protocol re-synced");
    }
}

impl Drop for CliProcess {
    fn drop(&mut self) {
        self.stderr_handle.abort();
        let _ = self.child.start_kill();
    }
}

/// Spawns the Claude Code CLI (`claude`) as a persistent child process using
/// `--input-format stream-json --output-format stream-json`. The CLI stays alive
/// across turns, maintaining conversation state internally. Messages are sent as
/// NDJSON on stdin with content arrays supporting text and image blocks. Responses
/// are NDJSON on stdout (`assistant` + `result` events per turn).
#[derive(Debug, serde::Serialize)]
pub struct ClaudeCodeProvider {
    command: PathBuf,
    model: ModelConfig,
    #[serde(skip)]
    name: String,
    /// Temp file holding MCP config JSON (auto-deleted on drop).
    #[serde(skip)]
    mcp_config_file: Option<NamedTempFile>,
    #[serde(skip)]
    cli_process: tokio::sync::OnceCell<Arc<tokio::sync::Mutex<CliProcess>>>,
    #[serde(skip)]
    pending_confirmations:
        Arc<tokio::sync::Mutex<HashMap<String, oneshot::Sender<PermissionConfirmation>>>>,
    #[serde(skip)]
    initial_mode: tokio::sync::Mutex<Option<GooseMode>>,
}

impl ClaudeCodeProvider {
    /// Build content blocks from the last user message only — the CLI maintains
    /// conversation context internally per session_id.
    fn last_user_content_blocks(&self, messages: &[Message]) -> Vec<Value> {
        let msgs = match messages.iter().rev().find(|m| m.role == Role::User) {
            Some(msg) => std::slice::from_ref(msg),
            None => messages,
        };
        let mut blocks: Vec<Value> = Vec::new();
        for message in msgs.iter().filter(|m| m.is_agent_visible()) {
            let prefix = match message.role {
                Role::User => "Human: ",
                Role::Assistant => "Assistant: ",
            };
            let mut text_parts = Vec::new();
            for content in &message.content {
                match content {
                    MessageContent::Text(t) => text_parts.push(t.text.clone()),
                    MessageContent::Image(img) => {
                        if !text_parts.is_empty() {
                            blocks.push(json!({"type":"text","text":format!("{}{}", prefix, text_parts.join("\n"))}));
                            text_parts.clear();
                        }
                        blocks.push(json!({"type":"image","source":{"type":"base64","media_type":img.mime_type,"data":img.data}}));
                    }
                    MessageContent::ToolRequest(req) => {
                        if let Ok(call) = &req.tool_call {
                            text_parts.push(format!("[tool_use: {} id={}]", call.name, req.id));
                        }
                    }
                    MessageContent::ToolResponse(resp) => {
                        if let Ok(result) = &resp.tool_result {
                            let text: String = result
                                .content
                                .iter()
                                .filter_map(|c| match &c.raw {
                                    rmcp::model::RawContent::Text(t) => Some(t.text.as_str()),
                                    _ => None,
                                })
                                .collect::<Vec<&str>>()
                                .join("\n");
                            text_parts.push(format!("[tool_result id={}] {}", resp.id, text));
                        }
                    }
                    _ => {}
                }
            }
            if !text_parts.is_empty() {
                blocks.push(
                    json!({"type":"text","text":format!("{}{}", prefix, text_parts.join("\n"))}),
                );
            }
        }
        blocks
    }

    fn build_stream_json_command(&self) -> Command {
        let mut cmd = Command::new(&self.command);
        configure_subprocess(&mut cmd);
        // Allow goose to run inside a Claude Code session.
        cmd.env_remove("CLAUDECODE");
        cmd.arg("--input-format")
            .arg("stream-json")
            .arg("--output-format")
            .arg("stream-json")
            .arg("--verbose")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        cmd
    }

    /// Returns true when the control protocol is enabled.
    fn apply_permission_flags(cmd: &mut Command) -> Result<bool, ProviderError> {
        let config = Config::global();
        let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);

        match goose_mode {
            GooseMode::Auto => {
                cmd.arg("--dangerously-skip-permissions");
                Ok(false)
            }
            GooseMode::SmartApprove | GooseMode::Approve => {
                cmd.arg("--permission-prompt-tool").arg("stdio");
                Ok(true)
            }
            GooseMode::Chat => Ok(false),
        }
    }

    async fn spawn_process(&self, filtered_system: &str) -> Result<CliProcess, ProviderError> {
        let mut cmd = self.build_stream_json_command();

        if let Some(f) = &self.mcp_config_file {
            cmd.arg("--mcp-config").arg(f.path());
            cmd.arg("--strict-mcp-config");
        }

        cmd.arg("--include-partial-messages")
            .arg("--system-prompt")
            .arg(filtered_system)
            .arg("--model")
            .arg(&self.model.model_name);

        let control_protocol_enabled = Self::apply_permission_flags(&mut cmd)?;

        let mut child = cmd.spawn().map_err(|e| {
            ProviderError::RequestFailed(format!(
                "Failed to spawn Claude CLI command '{:?}': {}.",
                self.command, e
            ))
        })?;

        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdin".to_string()))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdout".to_string()))?;

        let stderr = child.stderr.take();
        let stderr_handle = tokio::spawn(async move {
            let mut output = String::new();
            if let Some(mut stderr) = stderr {
                use tokio::io::AsyncReadExt;
                let _ = stderr.read_to_string(&mut output).await;
            }
            output
        });

        let mut process = CliProcess {
            child,
            stdin: Box::new(stdin),
            reader: BufReader::new(Box::new(stdout)),
            stderr_handle,
            current_model: self.model.model_name.clone(),
            log_model_update: false,
            next_request_id: 0,
            needs_drain: false,
        };

        if control_protocol_enabled {
            process
                .send_control_request(ControlRequestBody::Initialize)
                .await?;
        }

        Ok(process)
    }

    async fn get_or_init_process(
        &self,
        filtered_system: &str,
    ) -> Result<&Arc<tokio::sync::Mutex<CliProcess>>, ProviderError> {
        self.cli_process
            .get_or_try_init(|| async {
                Ok(Arc::new(tokio::sync::Mutex::new(
                    self.spawn_process(filtered_system).await?,
                )))
            })
            .await
    }
}

async fn exchange_control(
    stdin: &mut (impl AsyncWrite + Unpin),
    reader: &mut (impl AsyncBufRead + Unpin),
    request_id: &str,
    body: ControlRequestBody,
) -> Result<Option<Value>, ProviderError> {
    let label = body.label();
    let req = ControlRequest {
        msg_type: "control_request",
        request_id: request_id.to_string(),
        request: body,
    };
    let mut req_str = serde_json::to_string(&req).map_err(|e| {
        ProviderError::RequestFailed(format!("Failed to serialize {label} request: {e}"))
    })?;
    req_str.push('\n');
    stdin.write_all(req_str.as_bytes()).await.map_err(|e| {
        ProviderError::RequestFailed(format!("Failed to write {label} request: {e}"))
    })?;

    let mut line = String::new();
    loop {
        line.clear();
        match reader.read_line(&mut line).await {
            Ok(0) => {
                return Err(ProviderError::RequestFailed(format!(
                    "CLI process terminated while waiting for {label} response"
                )));
            }
            Ok(_) => {
                let trimmed = line.trim();
                if trimmed.is_empty() {
                    continue;
                }
                if let Ok(msg) = serde_json::from_str::<IncomingControlResponse>(trimmed) {
                    match msg.response {
                        IncomingControlResponseBody::Success {
                            request_id: ref rid,
                            response,
                        } if rid == request_id => return Ok(response),
                        IncomingControlResponseBody::Error {
                            request_id: ref rid,
                            error,
                        } if rid == request_id => {
                            return Err(ProviderError::RequestFailed(format!(
                                "{label} failed: {error}"
                            )));
                        }
                        _ => continue,
                    }
                }
            }
            Err(e) => {
                return Err(ProviderError::RequestFailed(format!(
                    "Failed to read {label} response: {e}"
                )));
            }
        }
    }
}

fn extract_model_aliases(response: Option<&Value>) -> Vec<String> {
    response
        .and_then(|v| v.get("models")?.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|m| m.get("value")?.as_str().map(String::from))
                .collect()
        })
        .unwrap_or_default()
}

fn build_stream_json_input(content_blocks: &[Value], session_id: &str) -> String {
    let msg = json!({"type":"user","session_id":session_id,"message":{"role":"user","content":content_blocks}});
    serde_json::to_string(&msg).expect("serializing JSON content blocks cannot fail")
}

fn claude_mcp_config_json(extensions: &[ExtensionConfig]) -> Option<String> {
    let mut mcp_servers = serde_json::Map::new();

    for extension in extensions {
        match extension {
            ExtensionConfig::StreamableHttp { uri, headers, .. } => {
                let key = extension.key();
                let mut config = serde_json::Map::new();
                config.insert("type".to_string(), json!("http"));
                config.insert("url".to_string(), json!(uri));
                if !headers.is_empty() {
                    config.insert("headers".to_string(), json!(headers));
                }
                mcp_servers.insert(key, Value::Object(config));
            }
            ExtensionConfig::Stdio {
                cmd, args, envs, ..
            } => {
                let key = extension.key();
                let mut config = serde_json::Map::new();
                config.insert("type".to_string(), json!("stdio"));
                config.insert("command".to_string(), json!(cmd));
                if !args.is_empty() {
                    config.insert("args".to_string(), json!(args));
                }
                let env_map = envs.get_env();
                if !env_map.is_empty() {
                    config.insert("env".to_string(), json!(env_map));
                }
                mcp_servers.insert(key, Value::Object(config));
            }
            ExtensionConfig::Sse { name, .. } => {
                tracing::debug!(name, "skipping SSE extension, migrate to streamable_http");
            }
            _ => {}
        }
    }

    if mcp_servers.is_empty() {
        return None;
    }

    serde_json::to_string(&json!({ "mcpServers": mcp_servers })).ok()
}

/// Write the MCP config JSON to a temp file with restricted permissions
/// so secrets (headers, env vars) are not leaked via process argv.
fn write_mcp_config_file(state_dir: &Path, json: &str) -> Result<NamedTempFile, anyhow::Error> {
    let dir = state_dir.join("claude-code");
    std::fs::create_dir_all(&dir)?;
    let prefix = format!("mcp-config-{}_", chrono::Utc::now().format("%Y%m%d"));
    let mut tmp = tempfile::Builder::new()
        .prefix(&prefix)
        .suffix(".json")
        .tempfile_in(&dir)?;
    tmp.write_all(json.as_bytes())?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        tmp.as_file()
            .set_permissions(std::fs::Permissions::from_mode(0o600))?;
    }
    Ok(tmp)
}

impl ProviderDef for ClaudeCodeProvider {
    type Provider = Self;

    fn metadata() -> ProviderMetadata {
        ProviderMetadata::new(
            CLAUDE_CODE_PROVIDER_NAME,
            "Claude Code CLI",
            "Requires claude CLI installed, no MCPs. Use Anthropic provider for full features.",
            CLAUDE_CODE_DEFAULT_MODEL,
            // Only a few agentic choices; fetched dynamically via fetch_supported_models.
            vec![],
            CLAUDE_CODE_DOC_URL,
            vec![ConfigKey::from_value_type::<ClaudeCodeCommand>(
                true, false, true,
            )],
        )
    }

    fn from_env(
        model: ModelConfig,
        extensions: Vec<ExtensionConfig>,
    ) -> BoxFuture<'static, Result<Self::Provider>> {
        Box::pin(async move {
            let config = crate::config::Config::global();
            let command: String = config.get_claude_code_command().unwrap_or_default().into();
            let resolved_command = SearchPaths::builder().with_npm().resolve(command)?;

            let mut resolved = Vec::with_capacity(extensions.len());
            for ext in extensions {
                resolved.push(ext.resolve(config).await?);
            }

            let mcp_config_file = claude_mcp_config_json(&resolved)
                .map(|json| write_mcp_config_file(&Paths::state_dir(), &json))
                .transpose()?;

            Ok(Self {
                command: resolved_command,
                model,
                name: CLAUDE_CODE_PROVIDER_NAME.to_string(),
                mcp_config_file,
                cli_process: tokio::sync::OnceCell::new(),
                pending_confirmations: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
                initial_mode: tokio::sync::Mutex::new(None),
            })
        })
    }
}

#[async_trait]
impl Provider for ClaudeCodeProvider {
    fn get_name(&self) -> &str {
        &self.name
    }

    fn get_model_config(&self) -> ModelConfig {
        self.model.clone()
    }

    async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
        // Uses a separate short-lived process because --system-prompt is a CLI-only
        // flag with no NDJSON equivalent. The persistent process needs it at spawn,
        // but it's unavailable during model listing.
        // See: https://code.claude.com/docs/en/cli-reference#system-prompt-flags
        let mut cmd = self.build_stream_json_command();
        let mut child = cmd.spawn().map_err(|e| {
            ProviderError::RequestFailed(format!("Failed to spawn CLI for model listing: {e}"))
        })?;

        let mut stdin = child
            .stdin
            .take()
            .ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdin".to_string()))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdout".to_string()))?;

        let mut reader = BufReader::new(stdout);
        let response = exchange_control(
            &mut stdin,
            &mut reader,
            "model_list",
            ControlRequestBody::Initialize,
        )
        .await;
        let _ = child.kill().await;
        Ok(extract_model_aliases(response.ok().flatten().as_ref()))
    }

    async fn update_mode(&self, _session_id: &str, mode: GooseMode) -> Result<(), ProviderError> {
        // Mode is baked into the subprocess at spawn; claude-acp replaces
        // this provider (#7801).
        let mut guard = self.initial_mode.lock().await;
        let current = *guard.get_or_insert(mode);
        if current != mode {
            return Err(ProviderError::RequestFailed(format!(
                "Mode change not supported: session is {current}, requested {mode}",
            )));
        }
        Ok(())
    }

    fn permission_routing(&self) -> PermissionRouting {
        PermissionRouting::ActionRequired
    }

    async fn handle_permission_confirmation(
        &self,
        request_id: &str,
        confirmation: &PermissionConfirmation,
    ) -> bool {
        let mut pending = self.pending_confirmations.lock().await;
        if let Some(tx) = pending.remove(request_id) {
            let _ = tx.send(confirmation.clone());
            return true;
        }
        false
    }

    async fn stream(
        &self,
        model_config: &ModelConfig,
        session_id: &str,
        system: &str,
        messages: &[Message],
        _tools: &[Tool],
    ) -> Result<MessageStream, ProviderError> {
        if super::cli_common::is_session_description_request(system) {
            let (message, usage) = super::cli_common::generate_simple_session_description(
                &model_config.model_name,
                messages,
            )?;
            return Ok(stream_from_single_message(message, usage));
        }

        let filtered_system = filter_extensions_from_system_prompt(system);
        let process_arc = Arc::clone(self.get_or_init_process(&filtered_system).await?);

        // Prepare the payload outside the lock — these don't need the process.
        let blocks = self.last_user_content_blocks(messages);
        let ndjson_line = build_stream_json_input(&blocks, session_id);
        let model_name = model_config.model_name.clone();
        let message_id = uuid::Uuid::new_v4().to_string();
        let pending_confirmations = Arc::clone(&self.pending_confirmations);

        Ok(Box::pin(try_stream! {
            // Single lock acquisition covers write-to-stdin and read-from-stdout,
            // eliminating the race window between the two.
            let mut process = process_arc.lock_owned().await;

            // Clean up pending permissions from a cancelled stream
            {
                let mut pending = pending_confirmations.lock().await;
                for (req_id, tx) in pending.drain() {
                    drop(tx);
                    let resp = ControlResponse::success(
                        req_id,
                        PermissionResponse::Deny { message: "Stream cancelled".to_string() },
                    );
                    let mut s = serde_json::to_string(&resp).map_err(|e| {
                        ProviderError::RequestFailed(format!("Failed to serialize cleanup deny response: {e}"))
                    })?;
                    s.push('\n');
                    let _ = process.stdin.write_all(s.as_bytes()).await;
                }
            }

            process.drain_pending_response().await;
            process.send_set_model(&model_name).await?;

            process
                .stdin
                .write_all(ndjson_line.as_bytes())
                .await
                .map_err(|e| {
                    ProviderError::RequestFailed(format!("Failed to write to stdin: {}", e))
                })?;
            process.stdin.write_all(b"\n").await.map_err(|e| {
                ProviderError::RequestFailed(format!("Failed to write newline to stdin: {}", e))
            })?;

            process.needs_drain = true;
            let mut line = String::new();
            let mut accumulated_usage = Usage::default();
            let mut stream_error: Option<ProviderError> = None;
            let stream_timestamp = chrono::Utc::now().timestamp();

            loop {
                line.clear();
                match process.reader.read_line(&mut line).await {
                    Ok(0) => {
                        process.needs_drain = false;
                        stream_error = Some(ProviderError::RequestFailed(
                            "Claude CLI process terminated unexpectedly".to_string(),
                        ));
                        break;
                    }
                    Ok(_) => {
                        let trimmed = line.trim();
                        if trimmed.is_empty() {
                            continue;
                        }

                        if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
                            match parsed.get("type").and_then(|t| t.as_str()) {
                                Some("stream_event") => {
                                    if let Some(event) = parsed.get("event") {
                                        match event.get("type").and_then(|t| t.as_str()) {
                                            Some("content_block_delta") => {
                                                if let Some(text) = event
                                                    .get("delta")
                                                    .filter(|d| {
                                                        d.get("type").and_then(|t| t.as_str())
                                                            == Some("text_delta")
                                                    })
                                                    .and_then(|d| d.get("text"))
                                                    .and_then(|t| t.as_str())
                                                {
                                                    if !text.is_empty() {
                                                        let mut partial_message = Message::new(
                                                            Role::Assistant,
                                                            stream_timestamp,
                                                            vec![MessageContent::text(text)],
                                                        );
                                                        partial_message.id =
                                                            Some(message_id.clone());
                                                        yield (Some(partial_message), None);
                                                    }
                                                }
                                            }
                                            Some("message_start") => {
                                                if let Some(usage_info) = event
                                                    .get("message")
                                                    .and_then(|m| m.get("usage"))
                                                {
                                                    let new = extract_usage_tokens(usage_info);
                                                    if let Some(i) = new.input_tokens {
                                                        accumulated_usage.input_tokens = Some(i);
                                                    }
                                                }
                                            }
                                            Some("message_delta") => {
                                                if let Some(usage_info) = event.get("usage") {
                                                    let new = extract_usage_tokens(usage_info);
                                                    if let Some(o) = new.output_tokens {
                                                        accumulated_usage.output_tokens = Some(o);
                                                    }
                                                }
                                            }
                                            _ => {}
                                        }
                                    }
                                }
                                Some("result") => {
                                    process.needs_drain = false;
                                    if let Some(usage_info) = parsed.get("usage") {
                                        let new = extract_usage_tokens(usage_info);
                                        accumulated_usage = Usage::new(
                                            new.input_tokens.or(accumulated_usage.input_tokens),
                                            new.output_tokens.or(accumulated_usage.output_tokens),
                                            None,
                                        );
                                    }
                                    break;
                                }
                                Some("error") => {
                                    process.needs_drain = false;
                                    stream_error = Some(error_from_event("Claude CLI", &parsed));
                                    break;
                                }
                                Some("control_request") => {
                                    if let Ok(IncomingControlRequest {
                                        request_id,
                                        request: IncomingRequestBody::CanUseTool { tool_name, input, tool_use_id },
                                    }) = serde_json::from_str::<IncomingControlRequest>(trimmed) {
                                        tracing::debug!(raw = %parsed, "can_use_tool control_request received");

                                        let (tx, rx) = oneshot::channel();
                                        pending_confirmations.lock().await.insert(request_id.clone(), tx);

                                        let action_msg = Message::assistant().with_action_required(
                                            request_id.clone(), tool_name, input.clone(), None,
                                        );
                                        yield (Some(action_msg), None);

                                        let confirmation = rx.await.unwrap_or(PermissionConfirmation {
                                            principal_type: PrincipalType::Tool,
                                            permission: Permission::Cancel,
                                        });
                                        pending_confirmations.lock().await.remove(&request_id);

                                        let perm_resp = match confirmation.permission {
                                            Permission::AlwaysAllow | Permission::AllowOnce => {
                                                PermissionResponse::Allow {
                                                    updated_input: input,
                                                    tool_use_id,
                                                }
                                            }
                                            _ => PermissionResponse::Deny {
                                                message: "User denied the tool call".to_string(),
                                            },
                                        };
                                        let resp = ControlResponse::success(request_id, perm_resp);
                                        let mut resp_str = serde_json::to_string(&resp).map_err(|e| {
                                            ProviderError::RequestFailed(format!("Failed to serialize permission response: {e}"))
                                        })?;
                                        tracing::debug!(json = %resp_str, "can_use_tool control_response sent");
                                        resp_str.push('\n');
                                        process.stdin.write_all(resp_str.as_bytes()).await.map_err(|e| {
                                            ProviderError::RequestFailed(format!("Failed to write permission response: {e}"))
                                        })?;
                                    }
                                }
                                Some("system") if process.log_model_update => {
                                    if let Some(resolved) = parsed.get("model").and_then(|m| m.as_str()) {
                                        tracing::debug!(
                                            from = %process.current_model,
                                            to = %resolved,
                                            "set_model resolved"
                                        );
                                    }
                                    process.log_model_update = false;
                                }
                                _ => {}
                            }
                        }
                    }
                    Err(e) => {
                        process.needs_drain = false;
                        stream_error = Some(ProviderError::RequestFailed(format!(
                            "Failed to read streaming output: {e}"
                        )));
                        break;
                    }
                }
            }

            if let Some(err) = stream_error {
                Err(err)?;
            }

            let provider_usage = ProviderUsage::new(model_name, accumulated_usage);
            yield (None, Some(provider_usage));
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agents::extension::Envs;
    use chrono::Utc;
    use goose_test_support::session::TEST_SESSION_ID;
    use serde_json::json;
    use std::collections::HashMap;
    use std::fs;
    use tempfile::tempdir;
    use test_case::test_case;

    #[test_case(
        json!({"input_tokens": 100, "output_tokens": 50}),
        Some(100), Some(50)
        ; "both_tokens"
    )]
    #[test_case(json!({"input_tokens": 100}), Some(100), None ; "input_only")]
    #[test_case(json!({}), None, None ; "empty_usage")]
    fn test_extract_usage_tokens(
        usage_json: Value,
        expected_input: Option<i32>,
        expected_output: Option<i32>,
    ) {
        let usage = extract_usage_tokens(&usage_json);
        assert_eq!(usage.input_tokens, expected_input);
        assert_eq!(usage.output_tokens, expected_output);
    }

    #[test_case(
        r#"{"type":"error","error":"context window exceeded"}"#,
        true
        ; "context_exceeded"
    )]
    #[test_case(
        r#"{"type":"error","error":"Model not supported"}"#,
        false
        ; "generic_error_from_event"
    )]
    #[test_case(r#"{"type":"error"}"#, false ; "missing_error_field")]
    fn test_error_from_event(line: &str, is_context_exceeded: bool) {
        let parsed: Value = serde_json::from_str(line).unwrap();
        let err = error_from_event("Claude CLI", &parsed);
        if is_context_exceeded {
            assert!(matches!(err, ProviderError::ContextLengthExceeded(_)));
        } else {
            assert!(matches!(err, ProviderError::RequestFailed(_)));
        }
    }

    /// (role, text, optional (image_data, mime_type))
    type MsgSpec<'a> = (&'a str, &'a str, Option<(&'a str, &'a str)>);

    fn build_messages(specs: &[MsgSpec]) -> Vec<Message> {
        specs
            .iter()
            .map(|(role, text, image)| {
                let role = if *role == "user" {
                    Role::User
                } else {
                    Role::Assistant
                };
                let mut msg = Message::new(role, 0, vec![]);
                if !text.is_empty() {
                    msg = Message::new(msg.role.clone(), 0, vec![MessageContent::text(*text)]);
                }
                if let Some((data, mime)) = image {
                    msg.content.push(MessageContent::image(*data, *mime));
                }
                msg
            })
            .collect()
    }

    #[test_case(
        build_messages(&[]),
        &[]
        ; "empty"
    )]
    #[test_case(
        build_messages(&[("user", "Hello", None)]),
        &[json!({"type":"text","text":"Human: Hello"})]
        ; "single_user"
    )]
    #[test_case(
        build_messages(&[("user", "Hello", None), ("assistant", "Hi there!", None)]),
        &[json!({"type":"text","text":"Human: Hello"})]
        ; "picks_last_user_ignores_assistant"
    )]
    #[test_case(
        build_messages(&[("user", "First", None), ("assistant", "Reply", None), ("user", "Second", None)]),
        &[json!({"type":"text","text":"Human: Second"})]
        ; "multi_turn_picks_last_user"
    )]
    #[test_case(
        build_messages(&[("user", "Describe this", Some(("base64data", "image/png")))]),
        &[json!({"type":"text","text":"Human: Describe this"}),
          json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":"base64data"}})]
        ; "user_with_image"
    )]
    #[test_case(
        build_messages(&[("user", "", Some(("iVBORw0KGgo", "image/png")))]),
        &[json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":"iVBORw0KGgo"}})]
        ; "image_only"
    )]
    #[test_case(
        vec![Message::new(Role::Assistant, 0, vec![
            MessageContent::tool_request("call_123", Ok(rmcp::model::CallToolRequestParams::new("developer__shell").with_arguments(serde_json::from_value(json!({"cmd": "ls"})).unwrap())))
        ])],
        &[json!({"type":"text","text":"Assistant: [tool_use: developer__shell id=call_123]"})]
        ; "tool_request_no_user_fallback"
    )]
    #[test_case(
        vec![Message::new(Role::User, 0, vec![
            MessageContent::tool_response("call_123", Ok(rmcp::model::CallToolResult::success(vec![rmcp::model::Content::text("file1.txt\nfile2.txt")])))
        ])],
        &[json!({"type":"text","text":"Human: [tool_result id=call_123] file1.txt\nfile2.txt"})]
        ; "tool_response"
    )]
    fn test_last_user_content_blocks(messages: Vec<Message>, expected: &[Value]) {
        let provider = make_provider();
        let blocks = provider.last_user_content_blocks(&messages);
        assert_eq!(blocks, expected);
    }

    #[test_case(
        &[json!({"type":"text","text":"Hello"})],
        json!({"type":"user","session_id":TEST_SESSION_ID,"message":{"role":"user","content":[{"type":"text","text":"Hello"}]}})
        ; "text_block"
    )]
    #[test_case(
        &[json!({"type":"text","text":"Look"}), json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"}})],
        json!({"type":"user","session_id":TEST_SESSION_ID,"message":{"role":"user","content":[{"type":"text","text":"Look"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"}}]}})
        ; "text_and_image_blocks"
    )]
    fn test_build_stream_json_input(blocks: &[Value], expected: Value) {
        let line = build_stream_json_input(blocks, TEST_SESSION_ID);
        let parsed: Value = serde_json::from_str(&line).unwrap();
        assert_eq!(parsed, expected);
    }

    #[test_case(
        Some(json!({"models":[{"value":"default","displayName":"Default"},{"value":"sonnet","displayName":"Sonnet"},{"value":"haiku","displayName":"Haiku"}]})),
        vec!["default".into(), "sonnet".into(), "haiku".into()]
        ; "success"
    )]
    #[test_case(
        Some(json!({"models":[{"value":"default","displayName":"Default"},{"value":null,"displayName":"Bad"}]})),
        vec!["default".into()]
        ; "filters_null_values"
    )]
    #[test_case(
        None,
        vec![]
        ; "none_input"
    )]
    #[test_case(
        Some(json!({"other":"data"})),
        vec![]
        ; "no_models_key"
    )]
    fn test_extract_model_aliases(response: Option<Value>, expected: Vec<String>) {
        assert_eq!(extract_model_aliases(response.as_ref()), expected);
    }

    #[test_case(
        vec![],
        None
        ; "empty_extensions_returns_none"
    )]
    #[test_case(
        vec![ExtensionConfig::Sse {
            name: "legacy".into(),
            description: String::new(),
            uri: Some("http://localhost/sse".into()),
        }],
        None
        ; "sse_only_returns_none"
    )]
    #[test_case(
        vec![ExtensionConfig::Stdio {
            name: "lookup".into(),
            description: String::new(),
            cmd: "node".into(),
            args: vec!["server.js".into()],
            envs: Envs::new([("API_KEY".into(), "secret".into())].into()),
            env_keys: vec![],
            timeout: None,
            bundled: Some(false),
            available_tools: vec![],
        }],
        Some(json!({ "mcpServers": {
            "lookup": {
                "type": "stdio",
                "command": "node",
                "args": ["server.js"],
                "env": { "API_KEY": "secret" }
            }
        }}))
        ; "stdio_converts_to_mcp_config_json"
    )]
    #[test_case(
        vec![ExtensionConfig::StreamableHttp {
            name: "lookup".into(),
            description: String::new(),
            uri: "http://localhost/mcp".into(),
            envs: Envs::default(),
            env_keys: vec![],
            headers: HashMap::from([("Authorization".into(), "Bearer token".into())]),
            timeout: None,
            bundled: Some(false),
            available_tools: vec![],
        }],
        Some(json!({ "mcpServers": {
            "lookup": {
                "type": "http",
                "url": "http://localhost/mcp",
                "headers": { "Authorization": "Bearer token" }
            }
        }}))
        ; "streamable_http_converts_to_mcp_config_json"
    )]
    #[test_case(
        vec![ExtensionConfig::StreamableHttp {
            name: "mcp_kiwi_com".into(),
            description: String::new(),
            uri: "https://mcp.kiwi.com".into(),
            envs: Envs::default(),
            env_keys: vec![],
            headers: HashMap::new(),
            timeout: None,
            bundled: None,
            available_tools: vec![],
        }],
        Some(json!({ "mcpServers": {
            "mcp_kiwi_com": {
                "type": "http",
                "url": "https://mcp.kiwi.com"
            }
        }}))
        ; "resolved_name_used_as_key"
    )]
    fn test_claude_mcp_config_json(extensions: Vec<ExtensionConfig>, expected: Option<Value>) {
        let result = claude_mcp_config_json(&extensions)
            .map(|json| serde_json::from_str::<Value>(&json).unwrap());
        assert_eq!(result, expected);
    }

    #[test]
    fn test_write_mcp_config_file() {
        let state_dir = tempdir().unwrap();
        let json = r#"{"mcpServers":{}}"#;

        let tmp = write_mcp_config_file(state_dir.path(), json).unwrap();

        assert_eq!(fs::read_to_string(tmp.path()).unwrap(), json);

        let norm_path = tmp.path().to_string_lossy().replace('\\', "/");
        let expected_prefix = format!("claude-code/mcp-config-{}_", Utc::now().format("%Y%m%d"));
        assert!(norm_path.contains(&expected_prefix));
        assert!(norm_path.ends_with(".json"));
    }

    #[test]
    fn test_write_mcp_config_file_invalid_state_dir() {
        assert!(write_mcp_config_file(Path::new("/dev/null"), "{}").is_err());
    }

    fn make_provider() -> ClaudeCodeProvider {
        ClaudeCodeProvider {
            command: PathBuf::from("claude"),
            model: ModelConfig::new(CLAUDE_CODE_DEFAULT_MODEL)
                .unwrap()
                .with_canonical_limits(CLAUDE_CODE_PROVIDER_NAME),
            name: "claude-code".to_string(),
            mcp_config_file: None,
            cli_process: tokio::sync::OnceCell::new(),
            pending_confirmations: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
            initial_mode: tokio::sync::Mutex::new(None),
        }
    }

    fn make_test_process(canned_stdout: &str) -> (CliProcess, tokio::io::DuplexStream) {
        let child = tokio::process::Command::new("true")
            .spawn()
            .expect("failed to spawn `true`");
        let (stdin_writer, stdin_reader) = tokio::io::duplex(1024);
        let process = CliProcess {
            child,
            stdin: Box::new(stdin_writer),
            reader: BufReader::new(Box::new(std::io::Cursor::new(
                canned_stdout.as_bytes().to_vec(),
            ))),
            stderr_handle: tokio::spawn(async { String::new() }),
            current_model: String::new(),
            log_model_update: false,
            next_request_id: 0,
            needs_drain: false,
        };
        (process, stdin_reader)
    }

    async fn stream_with_canned_stdout(
        canned_lines: &[&str],
    ) -> (ClaudeCodeProvider, MessageStream, tokio::io::DuplexStream) {
        let canned_stdout = canned_lines.join("\n");
        let (process, stdin_reader) = make_test_process(&canned_stdout);
        let provider = make_provider();
        let process_arc = Arc::new(tokio::sync::Mutex::new(process));
        provider.cli_process.set(process_arc).unwrap();

        let messages = vec![Message::user().with_text("test")];
        let stream = provider
            .stream(&provider.model, "test-session", "", &messages, &[])
            .await
            .unwrap();
        (provider, stream, stdin_reader)
    }

    async fn capture_stdin(
        provider: &ClaudeCodeProvider,
        mut reader: tokio::io::DuplexStream,
    ) -> String {
        use tokio::io::AsyncReadExt;
        provider.cli_process.get().unwrap().lock().await.stdin = Box::new(tokio::io::sink());
        let mut buf = Vec::new();
        reader.read_to_end(&mut buf).await.unwrap();
        String::from_utf8(buf).unwrap()
    }

    fn extract_permission_response(stdin_str: &str, request_id: &str) -> Value {
        let line = stdin_str
            .lines()
            .find(|l| l.contains(request_id) && l.contains("control_response"))
            .unwrap();
        let json: Value = serde_json::from_str(line).unwrap();
        json.pointer("/response/response").unwrap().clone()
    }

    #[test_case(
        &[r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#],
        Some("default"), "sonnet",
        Ok(()),
        "{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"sonnet\"}}\n"
        ; "default_to_sonnet"
    )]
    #[test_case(
        &[r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#],
        Some("sonnet"), "default",
        Ok(()),
        "{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"default\"}}\n"
        ; "sonnet_to_default"
    )]
    #[test_case(
        &[r#"{"type":"control_response","response":{"subtype":"error","request_id":"req_0","error":"bad model"}}"#],
        None, "bad",
        Err(ProviderError::RequestFailed("set_model failed: bad model".into())),
        "{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"bad\"}}\n"
        ; "failure"
    )]
    #[test_case(
        &[],
        Some("sonnet"), "sonnet",
        Ok(()), ""
        ; "skip_when_same_model"
    )]
    #[test_case(
        &[],
        None, "sonnet",
        Err(ProviderError::RequestFailed("CLI process terminated while waiting for set_model response".into())),
        "{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"sonnet\"}}\n"
        ; "eof"
    )]
    #[tokio::test]
    async fn test_send_set_model(
        lines: &[&str],
        initial_model: Option<&str>,
        target_model: &str,
        expected: Result<(), ProviderError>,
        expected_stdin: &str,
    ) {
        use tokio::io::AsyncReadExt;

        let stdout = lines.join("\n");
        let (mut process, mut stdin_reader) = make_test_process(&stdout);
        if let Some(m) = initial_model {
            process.current_model = m.to_string();
        }

        let result = process.send_set_model(target_model).await;
        process.stdin = Box::new(tokio::io::sink());
        let mut stdin_bytes = Vec::new();
        stdin_reader.read_to_end(&mut stdin_bytes).await.unwrap();

        assert_eq!(result, expected);
        if expected.is_ok() {
            assert_eq!(process.current_model, target_model);
        }
        assert_eq!(String::from_utf8(stdin_bytes).unwrap(), expected_stdin);
    }

    #[test_case(
        Permission::AllowOnce,
        json!({"behavior":"allow","updatedInput":{"path":"foo.txt","content":"hello"},"toolUseID":"tu_1"})
        ; "allow"
    )]
    #[test_case(
        Permission::DenyOnce,
        json!({"behavior":"deny","message":"User denied the tool call"})
        ; "deny"
    )]
    #[tokio::test]
    async fn test_can_use_tool(permission: Permission, expected_response: Value) {
        use futures::StreamExt;

        let (provider, mut stream, stdin_reader) = stream_with_canned_stdout(&[
            r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#,
            r#"{"type":"control_request","request_id":"perm_1","request":{"subtype":"can_use_tool","tool_name":"Write","input":{"path":"foo.txt","content":"hello"},"tool_use_id":"tu_1"}}"#,
            r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":5}}"#,
        ]).await;

        let (first_msg, _) = stream.next().await.unwrap().unwrap();
        let first_msg = first_msg.unwrap();
        let ar = first_msg
            .content
            .iter()
            .find_map(|c| c.as_action_required())
            .unwrap();
        match &ar.data {
            crate::conversation::message::ActionRequiredData::ToolConfirmation {
                id,
                tool_name,
                ..
            } => {
                assert_eq!(id, "perm_1");
                assert_eq!(tool_name, "Write");
            }
            _ => panic!("expected ToolConfirmation"),
        }

        let handled = provider
            .handle_permission_confirmation(
                "perm_1",
                &PermissionConfirmation {
                    principal_type: PrincipalType::Tool,
                    permission: permission.clone(),
                },
            )
            .await;
        assert!(handled);
        assert!(provider.pending_confirmations.lock().await.is_empty());

        while let Some(item) = stream.next().await {
            item.unwrap();
        }
        drop(stream);

        let stdin_str = capture_stdin(&provider, stdin_reader).await;
        let response_data = extract_permission_response(&stdin_str, "perm_1");
        assert_eq!(response_data, expected_response);
    }

    #[tokio::test]
    async fn test_can_use_tool_cancel_on_drop() {
        use futures::StreamExt;

        let (provider, mut stream, stdin_reader) = stream_with_canned_stdout(&[
            r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#,
            r#"{"type":"control_request","request_id":"perm_1","request":{"subtype":"can_use_tool","tool_name":"Write","input":{"path":"foo.txt"},"tool_use_id":"tu_1"}}"#,
            r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":5}}"#,
        ]).await;

        let pending = Arc::clone(&provider.pending_confirmations);

        let (first_msg, _) = stream.next().await.unwrap().unwrap();
        assert!(first_msg
            .unwrap()
            .content
            .iter()
            .any(|c| c.as_action_required().is_some()));

        let tx = pending.lock().await.remove("perm_1").unwrap();
        drop(tx);

        while let Some(item) = stream.next().await {
            item.unwrap();
        }
        drop(stream);

        let stdin_str = capture_stdin(&provider, stdin_reader).await;
        let response_data = extract_permission_response(&stdin_str, "perm_1");
        assert_eq!(
            response_data,
            json!({"behavior":"deny","message":"User denied the tool call"})
        );
    }

    #[tokio::test]
    async fn test_pending_permissions_cleaned_on_new_stream() {
        use futures::StreamExt;

        let canned_stdout = [
            r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#,
            r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":5}}"#,
        ]
        .join("\n");

        let (process, stdin_reader) = make_test_process(&canned_stdout);
        let provider = make_provider();
        let process_arc = Arc::new(tokio::sync::Mutex::new(process));
        provider.cli_process.set(process_arc).unwrap();

        let (tx, _rx) = oneshot::channel();
        provider
            .pending_confirmations
            .lock()
            .await
            .insert("stale_1".to_string(), tx);

        let messages = vec![Message::user().with_text("test")];
        let mut stream = provider
            .stream(&provider.model, "test-session", "", &messages, &[])
            .await
            .unwrap();

        while let Some(item) = stream.next().await {
            item.unwrap();
        }
        drop(stream);

        assert!(provider.pending_confirmations.lock().await.is_empty());

        let stdin_str = capture_stdin(&provider, stdin_reader).await;
        let response_data = extract_permission_response(&stdin_str, "stale_1");
        assert_eq!(response_data["behavior"], "deny");
    }
}