//! ACP agent loop with tools and session modes (Stage 3).
//!
//! Handlers:
//!
//! | ACP method            | Behaviour                                                   |
//! |-----------------------|-------------------------------------------------------------|
//! | `initialize`          | echo protocol version, advertise capabilities               |
//! | `session/new`         | mint id, register state, advertise [Default, Bypass] modes  |
//! | `session/prompt`      | tool-call loop: stream → dispatch tools → re-enter, repeat  |
//! | `session/cancel`      | fire the session's cancellation token                       |
//! | `session/set_mode`    | mutate the session's mode (gated vs. bypass-permissions)    |
//! | `session/set_model`   | switch the session's active model (endpoint:model selector) |
//! | (anything else)       | "not implemented yet" error                                 |
//!
//! Stage 5 flipped on image content. Stage 6 starts adding new wire
//! protocols (Anthropic /v1/messages first).

use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use agent_client_protocol::schema::{
    AgentCapabilities, CancelNotification, ContentBlock, InitializeRequest, InitializeResponse,
    ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, ModelId,
    ModelInfo as AcpModelInfo, NewSessionRequest, NewSessionResponse, PromptCapabilities,
    PromptRequest, PromptResponse, SessionCapabilities, SessionId, SessionInfo,
    SessionListCapabilities, SessionMode, SessionModeId, SessionModeState, SessionModelState,
    SessionNotification, SessionUpdate, SetSessionModeRequest, SetSessionModeResponse,
    SetSessionModelRequest, SetSessionModelResponse, StopReason, TextContent,
};
use agent_client_protocol::{Agent as AgentRole, Client, ConnectionTo, Dispatch, Stdio};
use futures::StreamExt;
use std::collections::BTreeMap;
use tokio_util::sync::CancellationToken;

use crate::compaction;
use crate::config::{Config, parse_model_selector};
use crate::prompt::build_system_prompt;
use crate::provider::{
    CompletionEvent, CompletionRequest, Message, MessageContent, Provider, Role, ToolCall,
};
use crate::session::{self, MODE_BYPASS, MODE_DEFAULT, MODE_PLAN, SessionState, SessionStore};
use crate::store::{self, PersistedSession};
use crate::tool_runner::{AcpClientOps, ToolCallEvent, dispatch_tool_call};
use crate::tools;

/// Maximum number of provider→tool→provider round-trips per
/// `session/prompt` request. Bound exists to keep a runaway model
/// from looping forever; the spec maps this to
/// [`StopReason::MaxTurnRequests`].
const MAX_TOOL_ROUNDS: usize = 25;

/// Public entry point. Wraps an `Arc<AgentInner>` so handlers can clone
/// it cheaply into every closure.
pub struct Agent {
    inner: Arc<AgentInner>,
}

struct AgentInner {
    /// Every successfully-built provider, indexed positionally. We look
    /// providers up by name (`endpoint:` prefix) rather than by index.
    providers: Vec<Arc<dyn Provider>>,
    /// Name of the endpoint used when a request omits the
    /// `endpoint:model` prefix.
    default_endpoint_name: String,
    /// Default model for the default endpoint, if configured. Required
    /// for Stage 2 because session/set_model lands in Stage 4 — a
    /// session with no model can't prompt anything.
    default_model: Option<String>,
    /// Per-endpoint `max_tokens` override. Looked up by endpoint
    /// name after resolution. `None` (or an absent entry) means the
    /// upstream picks its own default.
    max_tokens: std::collections::HashMap<String, u64>,
    /// Per-endpoint model context window in tokens. When set, the
    /// agent compacts history before each completion so the prompt
    /// fits inside `context_window - max_tokens - safety` tokens.
    /// Absent entry → no compaction (legacy behaviour).
    context_window: std::collections::HashMap<String, usize>,
    /// Aggregated list of selectable models across every configured
    /// endpoint, computed once at startup. With a single endpoint
    /// the model ids appear bare; with multiple endpoints every id
    /// carries the `endpoint:` prefix so the picker is unambiguous.
    /// Empty when every provider's `list_models` failed at startup —
    /// the dropdown then shows nothing and the session keeps using
    /// the configured `default_model`.
    available_models: Vec<AcpModelInfo>,
    sessions: SessionStore,
    system_prompt_path: Option<PathBuf>,
    /// Monotonic counter for minting session ids. The wire format is
    /// `hxa-{n}` — short, debuggable, and the protocol doesn't require
    /// UUIDs for session ids (it only requires them for message ids
    /// behind an unstable flag).
    next_session_id: AtomicU64,
}

impl Agent {
    /// Construct an agent from a validated [`Config`] and the providers
    /// that were successfully built for each endpoint.
    ///
    /// `async` because we call `Provider::list_models` on every
    /// provider up-front so the model-picker dropdown is populated
    /// from the very first `session/new`. Per-endpoint failure
    /// warns and skips rather than aborting startup — a single
    /// unreachable endpoint shouldn't take down the agent.
    pub async fn new(cfg: &Config, providers: Vec<Arc<dyn Provider>>) -> anyhow::Result<Self> {
        if providers.is_empty() {
            anyhow::bail!("no usable providers");
        }
        let default = cfg.default_endpoint();
        // The default endpoint's provider must have built successfully —
        // otherwise we can't honour `model = "bare-model-id"` requests.
        // (If only a non-default endpoint is usable, the operator should
        // promote it to `default_endpoint` in the TOML.)
        if !providers.iter().any(|p| p.name() == default.name) {
            anyhow::bail!(
                "default endpoint '{}' has no usable provider — check config",
                default.name
            );
        }
        let max_tokens = cfg
            .endpoints
            .iter()
            .filter_map(|ep| ep.max_tokens.map(|m| (ep.name.clone(), m)))
            .collect();
        let context_window = cfg
            .endpoints
            .iter()
            .filter_map(|ep| ep.context_window.map(|w| (ep.name.clone(), w)))
            .collect();
        let available_models = aggregate_models(&providers).await;
        tracing::info!(
            models = available_models.len(),
            endpoints = providers.len(),
            "model catalogue assembled"
        );
        Ok(Self {
            inner: Arc::new(AgentInner {
                providers,
                default_endpoint_name: default.name.clone(),
                default_model: default.default_model.clone(),
                max_tokens,
                context_window,
                available_models,
                sessions: session::new_store(),
                system_prompt_path: cfg.system_prompt_path.clone(),
                next_session_id: AtomicU64::new(1),
            }),
        })
    }

    /// Run the agent against an ACP transport (typically [`Stdio`]).
    /// Returns when the transport closes or a handler errors.
    pub async fn serve(self, transport: Stdio) -> agent_client_protocol::Result<()> {
        let inner = self.inner;
        AgentRole
            .builder()
            .name("helexa-acp")
            .on_receive_request(
                async move |req: InitializeRequest, responder, _cx| {
                    responder.respond(initialize_response(&req))
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_request(
                {
                    let inner = inner.clone();
                    async move |req: NewSessionRequest, responder, _cx| {
                        let result = handle_new_session(&inner, req).await;
                        match result {
                            Ok(resp) => responder.respond(resp),
                            Err(e) => responder.respond_with_internal_error(format!("{e:#}")),
                        }
                    }
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_request(
                {
                    let inner = inner.clone();
                    async move |req: LoadSessionRequest, responder, cx: ConnectionTo<Client>| {
                        let session_id = req.session_id.clone();
                        match handle_load_session(&inner, req).await {
                            Ok((resp, history)) => {
                                let send_result = responder.respond(resp);
                                // History replay happens off the
                                // dispatch loop so the load reply
                                // returns immediately. Zed receives
                                // the response, then sees a stream
                                // of session/update events that
                                // repopulate the chat panel.
                                let cx_clone = cx.clone();
                                let _ = cx.spawn(async move {
                                    replay_history(&cx_clone, &session_id, &history);
                                    Ok(())
                                });
                                send_result
                            }
                            Err(e) => responder.respond_with_internal_error(format!("{e:#}")),
                        }
                    }
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_request(
                async move |req: ListSessionsRequest, responder, _cx| match handle_list_sessions(
                    req,
                ) {
                    Ok(resp) => responder.respond(resp),
                    Err(e) => responder.respond_with_internal_error(format!("{e:#}")),
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_request(
                {
                    let inner = inner.clone();
                    async move |req: PromptRequest, responder, cx: ConnectionTo<Client>| {
                        spawn_prompt(inner.clone(), cx, req, responder)
                    }
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_request(
                {
                    let inner = inner.clone();
                    async move |req: SetSessionModeRequest, responder, _cx| {
                        match handle_set_session_mode(&inner, req).await {
                            Ok(()) => responder.respond(SetSessionModeResponse::new()),
                            Err(e) => responder.respond_with_internal_error(format!("{e:#}")),
                        }
                    }
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_request(
                {
                    let inner = inner.clone();
                    async move |req: SetSessionModelRequest, responder, _cx| {
                        match handle_set_session_model(&inner, req).await {
                            Ok(()) => responder.respond(SetSessionModelResponse::new()),
                            Err(e) => responder.respond_with_internal_error(format!("{e:#}")),
                        }
                    }
                },
                agent_client_protocol::on_receive_request!(),
            )
            .on_receive_notification(
                {
                    let inner = inner.clone();
                    async move |notif: CancelNotification, _cx: ConnectionTo<Client>| {
                        handle_cancel(&inner, notif).await;
                        Ok(())
                    }
                },
                agent_client_protocol::on_receive_notification!(),
            )
            .on_receive_dispatch(
                async move |message: Dispatch, cx: ConnectionTo<Client>| {
                    // `Dispatch` has three variants. For Request and
                    // Notification we want the "not implemented yet"
                    // error response. For *Response* we MUST forward
                    // the result to its awaiting `ResponseRouter` —
                    // otherwise our own outbound ACP calls
                    // (`fs/read_text_file`, `session/request_permission`,
                    // `terminal/*`, …) get their replies silently
                    // overwritten with whatever error we'd send a
                    // peer for an unknown method. That's how Stage 3
                    // tool dispatches were appearing as
                    // "Internal error: not implemented yet" results
                    // to the model.
                    match message {
                        Dispatch::Response(result, router) => router.respond_with_result(result),
                        other => {
                            tracing::warn!(
                                method = ?other.method(),
                                "unhandled ACP message"
                            );
                            other.respond_with_error(
                                agent_client_protocol::util::internal_error("not implemented yet"),
                                cx,
                            )
                        }
                    }
                },
                agent_client_protocol::on_receive_dispatch!(),
            )
            .connect_to(transport)
            .await
    }
}

fn initialize_response(req: &InitializeRequest) -> InitializeResponse {
    // Stage 5: image is on (Zed clipboard / drag-drop). Audio and
    // embedded resources flip on in later stages.
    let prompt_caps = PromptCapabilities::default().image(true);
    // Stage 3b: advertise both the top-level `load_session` flag and
    // the `session/list` sub-capability. Zed (and other ACP clients)
    // uses `session/list` to discover the session id that belongs to
    // a workspace before sending `session/load` — without it, the
    // client only knows how to mint new sessions and resume never
    // fires regardless of what's on disk.
    let session_caps =
        SessionCapabilities::default().list(Some(SessionListCapabilities::default()));
    InitializeResponse::new(req.protocol_version).agent_capabilities(
        AgentCapabilities::new()
            .prompt_capabilities(prompt_caps)
            .session_capabilities(session_caps)
            .load_session(true),
    )
}

async fn handle_new_session(
    inner: &AgentInner,
    req: NewSessionRequest,
) -> anyhow::Result<NewSessionResponse> {
    if !req.cwd.is_absolute() {
        anyhow::bail!("session cwd must be absolute, got {}", req.cwd.display());
    }
    let model_id = inner
        .default_model
        .clone()
        .ok_or_else(|| anyhow::anyhow!(
            "default endpoint '{}' has no default_model — set one in config or wait for Stage 4 set_model",
            inner.default_endpoint_name
        ))?;

    let n = inner.next_session_id.fetch_add(1, Ordering::Relaxed);
    let session_id = SessionId::new(format!("hxa-{n}"));
    let cwd_display = req.cwd.display().to_string();
    let log_model = model_id.clone();
    let state = SessionState::new(req.cwd, model_id);
    session::insert(&inner.sessions, session_id.clone(), state).await;

    tracing::info!(
        session_id = %session_id.0,
        model_id = %log_model,
        cwd = %cwd_display,
        "session created"
    );
    let resp = NewSessionResponse::new(session_id).modes(default_mode_state());
    let resp = match session_model_state(inner, &log_model) {
        Some(models) => resp.models(Some(models)),
        None => resp,
    };
    Ok(resp)
}

/// Rehydrate a session from disk.
///
/// Behaviour:
///
/// - Reads the persisted JSON from
///   `$XDG_DATA_HOME/helexa-acp/sessions/{id}.json`. Missing file →
///   error (Zed falls back to `session/new`).
/// - Overwrites the persisted `cwd` with the one the client just
///   sent. The user may have moved or symlinked the repo since
///   the session was first created; the *current* cwd is the
///   right place to root subsequent tool dispatches.
/// - Materialises an in-memory `SessionState` with the persisted
///   model + mode + history.
/// - Returns `LoadSessionResponse` carrying the same mode list as
///   `session/new`, plus the persisted `current_mode_id` so the
///   client renders the mode dropdown in the correct state.
async fn handle_load_session(
    inner: &AgentInner,
    req: LoadSessionRequest,
) -> anyhow::Result<(LoadSessionResponse, Vec<Message>)> {
    if !req.cwd.is_absolute() {
        anyhow::bail!("session cwd must be absolute, got {}", req.cwd.display());
    }
    let persisted = store::load(&req.session_id)?;
    // Snapshot the values we need for logging, the response, and
    // the post-load history replay before we move pieces of
    // `persisted` into `state`.
    let model_id = persisted.model_id.clone();
    let mode_id = persisted.mode_id.clone();
    let history_for_replay = persisted.history.clone();
    let history_turns = persisted.history.len();

    let mut state = SessionState::new(req.cwd.clone(), persisted.model_id);
    state.history = persisted.history;
    state.mode_id = SessionModeId::new(persisted.mode_id);
    session::insert(&inner.sessions, req.session_id.clone(), state).await;

    tracing::info!(
        session_id = %req.session_id.0,
        model_id = %model_id,
        mode = %mode_id,
        cwd = %req.cwd.display(),
        history_turns,
        "session loaded from disk"
    );

    let modes = SessionModeState::new(
        SessionModeId::new(mode_id),
        default_mode_state().available_modes,
    );
    let resp = LoadSessionResponse::new().modes(modes);
    let resp = match session_model_state(inner, &model_id) {
        Some(models) => resp.models(Some(models)),
        None => resp,
    };
    Ok((resp, history_for_replay))
}

/// Re-emit a session's persisted history as `session/update`
/// notifications so an ACP client (Zed) can render the prior chat
/// after a `session/load`. Without this, even a successful load
/// leaves the agent panel blank because Zed doesn't cache the
/// transcript client-side for custom agent_servers entries — that
/// caching only happens for first-party agents where Zed itself
/// owns the conversation state.
///
/// Mapping:
///
/// - `Role::User` text → `SessionUpdate::UserMessageChunk`
/// - `Role::Assistant` text → `SessionUpdate::AgentMessageChunk`
/// - `Role::Assistant` with tool calls → text chunk (if any) plus
///   one `ToolCall` event per call. We emit each with status =
///   `Completed` because the call already ran; the matching
///   `Role::Tool` result message is folded into the card's
///   content via a subsequent `ToolCallUpdate`.
/// - `Role::Tool` (tool result) → `ToolCallUpdate` carrying the
///   result text, keyed by `tool_call_id` so it lands on the
///   right card.
/// - `Role::System` → skipped; system prompts aren't rendered.
fn replay_history(cx: &ConnectionTo<Client>, session_id: &SessionId, history: &[Message]) {
    use agent_client_protocol::schema::{
        Content, ToolCall as AcpToolCall, ToolCallContent, ToolCallId, ToolCallStatus,
        ToolCallUpdate, ToolCallUpdateFields,
    };

    fn tool_kind_for(name: &str) -> agent_client_protocol::schema::ToolKind {
        use agent_client_protocol::schema::ToolKind;
        match name {
            "read_file" | "list_dir" => ToolKind::Read,
            "write_file" | "edit_file" => ToolKind::Edit,
            "bash" => ToolKind::Execute,
            _ => ToolKind::Other,
        }
    }
    fn title_for(name: &str, args_json: &str) -> String {
        match (
            name,
            serde_json::from_str::<serde_json::Value>(args_json).ok(),
        ) {
            ("read_file", Some(v)) => format!(
                "Read {}",
                v.get("path").and_then(|p| p.as_str()).unwrap_or("?")
            ),
            ("write_file", Some(v)) => format!(
                "Write {}",
                v.get("path").and_then(|p| p.as_str()).unwrap_or("?")
            ),
            ("edit_file", Some(v)) => format!(
                "Edit {}",
                v.get("path").and_then(|p| p.as_str()).unwrap_or("?")
            ),
            ("list_dir", Some(v)) => format!(
                "List {}",
                v.get("path").and_then(|p| p.as_str()).unwrap_or("?")
            ),
            ("bash", Some(v)) => {
                let cmd = v.get("command").and_then(|p| p.as_str()).unwrap_or("?");
                let snippet = if cmd.len() > 60 {
                    format!("{}…", &cmd[..60])
                } else {
                    cmd.to_string()
                };
                format!("Run: {snippet}")
            }
            (other, _) => format!("Tool: {other}"),
        }
    }

    let send = |update: SessionUpdate| {
        let notif = SessionNotification::new(session_id.clone(), update);
        if let Err(e) = cx.send_notification(notif) {
            tracing::warn!(
                error = %format!("{e:#}"),
                "replay: failed to forward history event"
            );
        }
    };

    let mut total_events: usize = 0;
    for msg in history {
        match (msg.role, &msg.content) {
            (Role::User, MessageContent::Text { text }) => {
                send(SessionUpdate::UserMessageChunk(text_chunk(text.clone())));
                total_events += 1;
            }
            (Role::User, MessageContent::MultiPart { parts }) => {
                // We can re-emit text parts as UserMessageChunks.
                // Images get a placeholder line so the user sees
                // *that* an image was attached; re-replaying the
                // image bytes themselves through ACP would require
                // a path round-trip we don't currently keep.
                for part in parts {
                    match part {
                        crate::provider::MessagePart::Text { text } => {
                            send(SessionUpdate::UserMessageChunk(text_chunk(text.clone())));
                            total_events += 1;
                        }
                        crate::provider::MessagePart::Image(img) => {
                            let label = match &img.uri {
                                Some(u) => format!("[image: {u}]"),
                                None => {
                                    format!("[image: {} ({} bytes)]", img.mime_type, img.data.len())
                                }
                            };
                            send(SessionUpdate::UserMessageChunk(text_chunk(label)));
                            total_events += 1;
                        }
                    }
                }
            }
            (Role::Assistant, MessageContent::Text { text }) => {
                send(SessionUpdate::AgentMessageChunk(text_chunk(text.clone())));
                total_events += 1;
            }
            (Role::Assistant, MessageContent::ToolCalls { text, calls }) => {
                if let Some(t) = text
                    && !t.is_empty()
                {
                    send(SessionUpdate::AgentMessageChunk(text_chunk(t.clone())));
                    total_events += 1;
                }
                for call in calls {
                    let raw_input = serde_json::from_str::<serde_json::Value>(&call.arguments)
                        .unwrap_or_else(|_| serde_json::Value::String(call.arguments.clone()));
                    let card = AcpToolCall::new(
                        ToolCallId::new(call.id.clone()),
                        title_for(&call.name, &call.arguments),
                    )
                    .kind(tool_kind_for(&call.name))
                    .status(ToolCallStatus::Completed)
                    .raw_input(raw_input);
                    send(SessionUpdate::ToolCall(card));
                    total_events += 1;
                }
            }
            (
                Role::Tool,
                MessageContent::ToolResult {
                    tool_call_id,
                    content,
                },
            ) => {
                let update = ToolCallUpdate::new(
                    ToolCallId::new(tool_call_id.clone()),
                    ToolCallUpdateFields::new()
                        .status(ToolCallStatus::Completed)
                        .content(vec![ToolCallContent::Content(Content::new(
                            ContentBlock::Text(TextContent::new(content.clone())),
                        ))]),
                );
                send(SessionUpdate::ToolCallUpdate(update));
                total_events += 1;
            }
            (Role::System, _) => {
                // System prompts aren't shown in the chat panel.
            }
            (role, content) => {
                tracing::debug!(
                    ?role,
                    ?content,
                    "replay: unrecognised (role, content) shape; skipping"
                );
            }
        }
    }

    tracing::info!(
        session_id = %session_id.0,
        events = total_events,
        history_turns = history.len(),
        "session history replayed to client"
    );
}

/// Enumerate persisted sessions for the `session/list` ACP method.
///
/// Zed calls this on workspace open to find the session belonging
/// to the cwd it's reopening — without it, even though `session/load`
/// works, the client has no way to discover the session_id and
/// always falls back to `session/new`. That's exactly the
/// "history didn't survive the restart" symptom.
///
/// Cursor pagination from the request is accepted but ignored:
/// helexa-acp's session counts are too small to need it. We always
/// return the whole filtered list with `next_cursor = None`.
fn handle_list_sessions(req: ListSessionsRequest) -> anyhow::Result<ListSessionsResponse> {
    let sessions = store::list(req.cwd.as_deref())?;
    let infos: Vec<SessionInfo> = sessions
        .into_iter()
        .map(|s| {
            let mut info = SessionInfo::new(SessionId::new(s.session_id), s.cwd);
            info = info.title(derive_session_title(&s.history));
            info = info.updated_at(store::unix_to_iso8601(s.updated_at));
            info
        })
        .collect();
    tracing::info!(
        cwd = ?req.cwd,
        count = infos.len(),
        "session/list responded"
    );
    Ok(ListSessionsResponse::new(infos))
}

/// Best-effort human-readable title for a session, derived from the
/// first user turn's text (truncated to ~60 chars). Empty string
/// becomes `None` so Zed can fall back to its own placeholder.
fn derive_session_title(history: &[Message]) -> Option<String> {
    use crate::provider::MessagePart;
    history
        .iter()
        .find_map(|msg| match (msg.role, &msg.content) {
            (Role::User, MessageContent::Text { text }) => Some(text.clone()),
            (Role::User, MessageContent::MultiPart { parts }) => parts.iter().find_map(|p| {
                if let MessagePart::Text { text } = p {
                    Some(text.clone())
                } else {
                    None
                }
            }),
            _ => None,
        })
        .map(|s| {
            let trimmed = s.trim();
            if trimmed.chars().count() > 60 {
                let prefix: String = trimmed.chars().take(60).collect();
                format!("{prefix}…")
            } else {
                trimmed.to_string()
            }
        })
        .filter(|s| !s.is_empty())
}

/// Build the model catalogue advertised in `NewSessionResponse.models`
/// (and the resume equivalent). Walks every provider, calls
/// `list_models`, and prefixes ids with `endpoint:` when the user has
/// more than one endpoint configured. A failing endpoint logs and
/// contributes nothing — losing one endpoint must not blank the whole
/// dropdown.
async fn aggregate_models(providers: &[Arc<dyn Provider>]) -> Vec<AcpModelInfo> {
    let multi_endpoint = providers.len() > 1;
    let mut out: Vec<AcpModelInfo> = Vec::new();
    for provider in providers {
        let endpoint = provider.name().to_string();
        match provider.list_models().await {
            Ok(models) => {
                tracing::info!(
                    endpoint = %endpoint,
                    count = models.len(),
                    "fetched models from endpoint"
                );
                for m in models {
                    let id = if multi_endpoint {
                        format!("{endpoint}:{}", m.id)
                    } else {
                        m.id.clone()
                    };
                    let display = m.display_name.unwrap_or_else(|| m.id.clone());
                    let info = AcpModelInfo::new(ModelId::new(id), display)
                        .description(Some(format!("endpoint: {endpoint}")));
                    out.push(info);
                }
            }
            Err(e) => {
                tracing::warn!(
                    endpoint = %endpoint,
                    error = %format!("{e:#}"),
                    "list_models failed; this endpoint's models won't appear in the picker"
                );
            }
        }
    }
    out
}

/// Build the `SessionModelState` that Zed renders as the
/// model-picker dropdown. The current model id is exactly what
/// the session is using right now (already in `endpoint:model`
/// form if it was set that way). Returns `None` when the
/// catalogue is empty — no point showing an empty dropdown.
fn session_model_state(inner: &AgentInner, current: &str) -> Option<SessionModelState> {
    if inner.available_models.is_empty() {
        return None;
    }
    Some(SessionModelState::new(
        ModelId::new(current.to_string()),
        inner.available_models.clone(),
    ))
}

async fn handle_set_session_model(
    inner: &AgentInner,
    req: SetSessionModelRequest,
) -> anyhow::Result<()> {
    let Some(state) = session::get(&inner.sessions, &req.session_id).await else {
        anyhow::bail!("unknown session id {}", req.session_id.0);
    };
    let target = req.model_id.0.as_ref().to_string();
    // Validate the requested model id resolves to a configured
    // provider. We don't require it to appear in `available_models`
    // because the catalogue may be stale (endpoint added a model
    // after startup) and rejecting unknown ids would be too rigid.
    // Provider lookup is the actual source of truth.
    let (_, _) = resolve_provider(&inner.providers, &inner.default_endpoint_name, &target)
        .map_err(|e| anyhow::anyhow!("set_session_model: {e:#}"))?;
    // Persist the new model id on the session under the mutex,
    // then snapshot for disk persistence outside the lock.
    let snapshot = {
        let mut s = state.lock().await;
        s.model_id = target.clone();
        PersistedSession {
            session_id: req.session_id.0.as_ref().to_string(),
            cwd: s.cwd.clone(),
            model_id: s.model_id.clone(),
            mode_id: s.mode_id.0.as_ref().to_string(),
            history: s.history.clone(),
            created_at: store::now_secs(),
            updated_at: store::now_secs(),
        }
    };
    if let Err(e) = store::save(&snapshot) {
        tracing::warn!(
            session_id = %req.session_id.0,
            error = %format!("{e:#}"),
            "session persist after set_model failed; on-disk model id stays stale"
        );
    }
    tracing::info!(
        session_id = %req.session_id.0,
        model_id = %target,
        "session model changed"
    );
    Ok(())
}

/// The three modes every Stage 3 session advertises:
///
/// - **Default** — writes / bash prompt the user.
/// - **Bypass Permissions** — auto-allow.
/// - **Plan** — read-and-plan-only. Writes are restricted to a
///   per-project plan directory under `$XDG_DATA_HOME/helexa-acp/plans/`
///   and bash is disabled. Designed for "draft the implementation
///   plan, then I'll review and let you execute" flows.
fn default_mode_state() -> SessionModeState {
    SessionModeState::new(
        SessionModeId::new(MODE_DEFAULT),
        vec![
            SessionMode::new(SessionModeId::new(MODE_DEFAULT), "Default")
                .description("Prompt for permission before writes or shell commands."),
            SessionMode::new(SessionModeId::new(MODE_BYPASS), "Bypass Permissions")
                .description("Auto-allow all tool calls. Use with care."),
            SessionMode::new(SessionModeId::new(MODE_PLAN), "Plan")
                .description("Write plans to the plan directory; no shell, no writes outside it."),
        ],
    )
}

async fn handle_set_session_mode(
    inner: &AgentInner,
    req: SetSessionModeRequest,
) -> anyhow::Result<()> {
    let Some(state) = session::get(&inner.sessions, &req.session_id).await else {
        anyhow::bail!("unknown session id {}", req.session_id.0);
    };
    let accepted = matches!(
        req.mode_id.0.as_ref(),
        MODE_DEFAULT | MODE_BYPASS | MODE_PLAN
    );
    if !accepted {
        anyhow::bail!(
            "unknown mode '{}' — must be one of: {}, {}, {}",
            req.mode_id.0,
            MODE_DEFAULT,
            MODE_BYPASS,
            MODE_PLAN
        );
    }
    state.lock().await.mode_id = req.mode_id.clone();
    tracing::info!(
        session_id = %req.session_id.0,
        mode = %req.mode_id.0,
        "session mode changed"
    );
    Ok(())
}

async fn handle_cancel(inner: &AgentInner, notif: CancelNotification) {
    let Some(state) = session::get(&inner.sessions, &notif.session_id).await else {
        tracing::debug!(session_id = %notif.session_id.0, "cancel for unknown session, ignoring");
        return;
    };
    let cancel = state.lock().await.cancel.clone();
    tracing::info!(session_id = %notif.session_id.0, "cancellation requested");
    cancel.cancel();
}

/// Kick the prompt off on a spawned task so the event loop is free to
/// dispatch the matching `session/cancel`. The handler itself returns
/// `Ok(())` immediately (= `Handled::Yes`); the spawned task is what
/// eventually consumes `responder`.
fn spawn_prompt(
    inner: Arc<AgentInner>,
    cx: ConnectionTo<Client>,
    req: PromptRequest,
    responder: agent_client_protocol::Responder<PromptResponse>,
) -> agent_client_protocol::Result<()> {
    let task_cx = cx.clone();
    cx.spawn(async move {
        if let Err(e) = drive_prompt(inner, task_cx, req, responder).await {
            // `drive_prompt` already consumed the responder on the
            // error paths it produces; this branch only fires if the
            // task itself errored before reaching responder.respond.
            // Log and swallow — propagating the error would tear down
            // the whole connection, which is too violent for one
            // failed prompt.
            tracing::error!(error = %format!("{e:#}"), "prompt task failed");
        }
        Ok(())
    })?;
    Ok(())
}

async fn drive_prompt(
    inner: Arc<AgentInner>,
    cx: ConnectionTo<Client>,
    req: PromptRequest,
    responder: agent_client_protocol::Responder<PromptResponse>,
) -> anyhow::Result<()> {
    let session_id = req.session_id.clone();
    let Some(session_arc) = session::get(&inner.sessions, &session_id).await else {
        let _ =
            responder.respond_with_internal_error(format!("unknown session id {}", session_id.0));
        return Ok(());
    };

    // Snapshot the inputs under the session lock, then drop the lock
    // before any `await` that touches the network. `mode_id` is
    // refreshed at the top of every round (the user can toggle modes
    // mid-turn and we want the next round's streaming + tool gating
    // to reflect that).
    let (existing_history, model_id, cwd, cancel, mut mode_id) = {
        let mut state = session_arc.lock().await;
        // Fire the session's current cancel before installing a new
        // one. If a previous prompt task is still in-flight (model
        // stalled mid-stream, a long-running bash, a wedged ACP
        // roundtrip), this lets it observe is_cancelled() at the
        // next .await and unwind cleanly — instead of two tasks
        // racing each other to mutate session.history and to
        // persist the same file.
        state.cancel.cancel();
        let cancel = CancellationToken::new();
        state.cancel = cancel.clone();
        let user_content = flatten_prompt(&req.prompt);
        state.history.push(Message {
            role: Role::User,
            content: user_content,
        });
        (
            state.history.clone(),
            state.model_id.clone(),
            state.cwd.clone(),
            cancel,
            state.mode_id.clone(),
        )
    };

    let tool_specs = tools::all_tools();
    // Plan-mode write target. Resolved once because the cwd doesn't
    // change for a session's lifetime; the directory is created
    // lazily by the runtime when a write lands inside it.
    let plan_dir = store::plan_dir_for(&cwd);

    let (provider, local_model) =
        match resolve_provider(&inner.providers, &inner.default_endpoint_name, &model_id) {
            Ok(pair) => pair,
            Err(e) => {
                let _ = responder.respond_with_internal_error(format!("{e:#}"));
                return Ok(());
            }
        };

    tracing::info!(
        session_id = %session_id.0,
        endpoint = %provider.name(),
        model = %local_model,
        mode = %mode_id.0,
        history_turns = existing_history.len(),
        "sending prompt upstream"
    );

    let ops = AcpClientOps::new(cx.clone());

    // `messages` is the rolling conversation we send to the provider
    // each round. Slot 0 is the system prompt — rebuilt at the top
    // of every round so a mid-turn mode toggle takes effect. We seed
    // a placeholder here and overwrite it on the first iteration.
    let mut messages: Vec<Message> = Vec::with_capacity(existing_history.len() + 1);
    messages.push(Message {
        role: Role::System,
        content: MessageContent::Text {
            text: String::new(),
        },
    });
    messages.extend(existing_history);

    // Buffer for turns produced this round. Flushed into
    // session.history *and* persisted at the end of every iteration
    // (and once more after the loop). Per-round persistence means
    // a stall later in the conversation doesn't lose earlier rounds.
    let mut new_turns: Vec<Message> = Vec::new();
    // Monotonic counter for synthetic ids assigned to unparseable
    // <tool_call> blocks across all rounds of this prompt.
    let mut next_malformed_index: usize = 0;

    let mut stop_reason = StopReason::EndTurn;

    for round in 0..MAX_TOOL_ROUNDS {
        if cancel.is_cancelled() {
            stop_reason = StopReason::Cancelled;
            break;
        }

        // Refresh mode + rebuild system prompt at the top of every
        // round. Cheap (one mutex acquisition + one string build);
        // the win is that if the user flips the mode dropdown
        // mid-turn — particularly the Plan ↔ Bypass transitions
        // the plan-mode menu invites them to make — the new mode
        // gates both this round's streaming *and* its tool
        // dispatch.
        mode_id = session_arc.lock().await.mode_id.clone();
        let system_prompt = build_system_prompt(
            &cwd,
            inner.system_prompt_path.as_deref(),
            &tool_specs,
            &mode_id,
            plan_dir.as_deref(),
        )
        .map_err(|e| anyhow::anyhow!("build system prompt: {e:#}"))?;
        messages[0] = Message {
            role: Role::System,
            content: MessageContent::Text {
                text: system_prompt,
            },
        };

        tracing::info!(
            session_id = %session_id.0,
            round = round + 1,
            of = MAX_TOOL_ROUNDS,
            mode = %mode_id.0,
            history_turns = messages.len(),
            "prompt round: streaming"
        );

        // Project history into the model's context window when the
        // endpoint advertises one. Compaction is a per-request
        // *projection* — `messages` (and the persisted session
        // history downstream) stay intact; only what we send
        // upstream shrinks. Without this, a 32 K Qwen3 dies after
        // the first few `read_file` results pile up in history.
        let provider_max_tokens = inner.max_tokens.get(provider.name()).copied();
        let messages_for_provider = match inner.context_window.get(provider.name()).copied() {
            Some(ctx) => {
                let budget = prompt_budget(ctx, provider_max_tokens);
                let (compacted, stats) = compaction::compact_to_budget(&messages, budget);
                if stats.elided_messages > 0 {
                    tracing::info!(
                        session_id = %session_id.0,
                        round = round + 1,
                        context_window = ctx,
                        budget,
                        original_tokens = stats.original_tokens,
                        final_tokens = stats.final_tokens,
                        elided = stats.elided_messages,
                        "context compaction applied"
                    );
                }
                compacted
            }
            None => messages.clone(),
        };

        // Tool descriptions reach the model via the Qwen3 `# Tools`
        // block in the system prompt, not via the OpenAI `tools`
        // request field — cortex/neuron pass that field through to
        // the encoder unread, and including it would double-describe
        // tools once a strict-OpenAI backend lands. Leave empty.
        let completion_req = CompletionRequest {
            model: local_model.clone(),
            messages: messages_for_provider,
            tools: vec![],
            temperature: None,
            top_p: None,
            max_tokens: provider_max_tokens,
        };

        let mut stream = match provider.complete(completion_req, cancel.clone()).await {
            Ok(s) => s,
            Err(e) => {
                let _ = responder
                    .respond_with_internal_error(format!("{} complete: {e:#}", provider.name()));
                return Ok(());
            }
        };

        let mut assistant_text = String::new();
        let mut finish_reason: Option<String> = None;
        // `BTreeMap` keyed by the provider's tool-call index keeps
        // insertion order while allowing arg deltas to mutate any
        // bucket — `ToolCallStart` may arrive interleaved with
        // `ToolCallArgsDelta` for different indices.
        let mut tool_buckets: BTreeMap<usize, ToolCallBucket> = BTreeMap::new();
        // <tool_call> blocks whose JSON couldn't be parsed even with
        // qwen3's repair pass. We surface each as a Failed
        // ToolCall card and feed a synthetic error result back to
        // the model so it can retry on the next round.
        let mut malformed_calls: Vec<String> = Vec::new();

        while let Some(event) = stream.next().await {
            let event = match event {
                Ok(e) => e,
                Err(e) => {
                    tracing::warn!(error = %format!("{e:#}"), "stream error; ending round");
                    break;
                }
            };
            match event {
                CompletionEvent::TextDelta(t) => {
                    assistant_text.push_str(&t);
                    send_chunk(
                        &cx,
                        &session_id,
                        SessionUpdate::AgentMessageChunk(text_chunk(t)),
                    );
                }
                CompletionEvent::ReasoningDelta(t) => {
                    send_chunk(
                        &cx,
                        &session_id,
                        SessionUpdate::AgentThoughtChunk(text_chunk(t)),
                    );
                }
                CompletionEvent::ToolCallStart { index, id, name } => {
                    tool_buckets.insert(
                        index,
                        ToolCallBucket {
                            id,
                            name,
                            arguments: String::new(),
                        },
                    );
                }
                CompletionEvent::ToolCallArgsDelta { index, args_delta } => {
                    tool_buckets
                        .entry(index)
                        .or_default()
                        .arguments
                        .push_str(&args_delta);
                }
                CompletionEvent::MalformedToolCall { raw } => {
                    malformed_calls.push(raw);
                }
                CompletionEvent::Finish { reason } => finish_reason = reason,
                CompletionEvent::Usage(_) => {}
            }
        }

        if cancel.is_cancelled() {
            stop_reason = StopReason::Cancelled;
            // Persist any partial text so the next turn has context.
            if !assistant_text.is_empty() {
                new_turns.push(Message {
                    role: Role::Assistant,
                    content: MessageContent::Text {
                        text: assistant_text,
                    },
                });
            }
            break;
        }

        // Recovery pass before deciding "is there work to do?".
        // For each malformed body, try shape-based inference
        // against the tool catalogue (handles the "model emitted
        // `arguments` but forgot `name`" case). Successes get
        // promoted to real tool buckets; failures stay in
        // `malformed_calls` for the Failed-card path below.
        malformed_calls.retain(|raw| match try_repair_missing_name(raw) {
            Some((name, args_json)) => {
                let idx = tool_buckets
                    .keys()
                    .max()
                    .copied()
                    .map(|m| m + 1)
                    .unwrap_or(0);
                tracing::debug!(
                    inferred_name = %name,
                    index = idx,
                    "qwen3: recovered missing-name tool call via shape inference"
                );
                tool_buckets.insert(
                    idx,
                    ToolCallBucket {
                        id: format!("call_recovered_{idx}"),
                        name,
                        arguments: args_json,
                    },
                );
                false
            }
            None => true,
        });

        let has_tool_calls = !tool_buckets.is_empty();
        let has_malformed = !malformed_calls.is_empty();

        if !has_tool_calls && !has_malformed {
            // Terminal turn: just text. Save and finish.
            if !assistant_text.is_empty() {
                new_turns.push(Message {
                    role: Role::Assistant,
                    content: MessageContent::Text {
                        text: assistant_text,
                    },
                });
            }
            stop_reason = map_finish_reason(finish_reason.as_deref());
            break;
        }

        // Assistant turn carrying any successfully-parsed tool calls
        // (malformed ones are handled separately so each gets its
        // own Failed card with its raw body intact).
        let calls: Vec<ToolCall> = tool_buckets
            .values()
            .map(|b| ToolCall {
                id: b.id.clone(),
                name: b.name.clone(),
                arguments: b.arguments.clone(),
            })
            .collect();
        if has_tool_calls || !assistant_text.is_empty() {
            let assistant_turn = Message {
                role: Role::Assistant,
                content: if has_tool_calls {
                    MessageContent::ToolCalls {
                        text: (!assistant_text.is_empty()).then_some(assistant_text),
                        calls,
                    }
                } else {
                    MessageContent::Text {
                        text: assistant_text,
                    }
                },
            };
            new_turns.push(assistant_turn.clone());
            messages.push(assistant_turn);
        }

        // Dispatch every tool call sequentially. Parallelism is
        // tempting but would require Zed to handle interleaved
        // permission prompts; serial is friendlier.
        for bucket in tool_buckets.into_values() {
            if cancel.is_cancelled() {
                stop_reason = StopReason::Cancelled;
                break;
            }
            let event = ToolCallEvent {
                id: bucket.id,
                name: bucket.name,
                arguments: bucket.arguments,
            };
            tracing::info!(
                session_id = %session_id.0,
                tool = %event.name,
                tool_call_id = %event.id,
                "dispatch tool"
            );
            let result =
                dispatch_tool_call(&ops, &session_id, &mode_id, &cwd, event, &cancel).await;
            tracing::info!(
                session_id = %session_id.0,
                tool_call_id = %result.tool_call_id,
                is_error = result.is_error,
                "dispatch tool complete"
            );
            let result_turn = Message {
                role: Role::Tool,
                content: MessageContent::ToolResult {
                    tool_call_id: result.tool_call_id,
                    content: result.content,
                },
            };
            new_turns.push(result_turn.clone());
            messages.push(result_turn);
        }

        // Handle malformed calls last — each becomes a Failed
        // SessionUpdate::ToolCall card (so Zed renders structured
        // failure UI instead of dumping raw JSON inline) plus a
        // synthetic tool-result message so the model gets concrete
        // feedback for self-correction on the next round.
        for raw in malformed_calls.drain(..) {
            if cancel.is_cancelled() {
                stop_reason = StopReason::Cancelled;
                break;
            }
            let synthetic_id = next_synthetic_id(&mut next_malformed_index);
            emit_malformed_tool_card(&cx, &session_id, &synthetic_id, &raw);
            let (call_turn, result_turn) = synthesize_malformed_history(&synthetic_id, &raw);
            new_turns.push(call_turn.clone());
            messages.push(call_turn);
            new_turns.push(result_turn.clone());
            messages.push(result_turn);
        }

        if cancel.is_cancelled() {
            stop_reason = StopReason::Cancelled;
            break;
        }

        if round + 1 == MAX_TOOL_ROUNDS {
            tracing::warn!(
                session_id = %session_id.0,
                rounds = MAX_TOOL_ROUNDS,
                "hit MAX_TOOL_ROUNDS, returning MaxTurnRequests"
            );
            stop_reason = StopReason::MaxTurnRequests;
        }

        // Per-round flush: push this round's turns into the in-memory
        // history and persist to disk. If the model stalls in a later
        // round (long bash, upstream SSE that never finishes, etc.)
        // earlier rounds still survive a binary restart.
        if !new_turns.is_empty() {
            let drained = std::mem::take(&mut new_turns);
            tracing::info!(
                session_id = %session_id.0,
                round = round + 1,
                turns = drained.len(),
                "prompt round complete; persisting"
            );
            extend_and_persist(&session_arc, &session_id, drained).await;
        }
    }

    // Final flush for whatever the break paths above left behind.
    // No-op when the per-round flush already drained new_turns.
    if !new_turns.is_empty() {
        extend_and_persist(&session_arc, &session_id, new_turns).await;
    }

    tracing::info!(
        session_id = %session_id.0,
        ?stop_reason,
        "prompt complete"
    );
    let _ = responder.respond(PromptResponse::new(stop_reason));
    Ok(())
}

/// Push `new_turns` into the session's in-memory history under the
/// session lock, then snapshot the full state and write it to disk
/// *outside* the lock. Used by `drive_prompt` at the end of every
/// tool-call round (so partial progress survives a stall) and once
/// more after the loop (catching any turns the break paths left
/// behind).
///
/// Persistence failures log a warning and don't propagate — losing
/// a save shouldn't tear down a live conversation.
async fn extend_and_persist(
    session_arc: &Arc<tokio::sync::Mutex<SessionState>>,
    session_id: &SessionId,
    new_turns: Vec<Message>,
) {
    let snapshot = {
        let mut state = session_arc.lock().await;
        state.history.extend(new_turns);
        PersistedSession {
            session_id: session_id.0.as_ref().to_string(),
            cwd: state.cwd.clone(),
            model_id: state.model_id.clone(),
            mode_id: state.mode_id.0.as_ref().to_string(),
            history: state.history.clone(),
            // `created_at` ought to be preserved across saves —
            // currently SessionState doesn't carry it, so every
            // save refreshes both timestamps. Acceptable for
            // resume; future work: thread `created_at` through.
            created_at: store::now_secs(),
            updated_at: store::now_secs(),
        }
    };
    if let Err(e) = store::save(&snapshot) {
        tracing::warn!(
            session_id = %session_id.0,
            error = %format!("{e:#}"),
            "session/persist failed; resume from disk will miss this round"
        );
    }
}

/// Accumulator for one streamed tool call: the OpenAI wire format
/// sends `id` + `name` once (in the first chunk for that index) and
/// then argument bytes piecemeal. We gather them all before
/// dispatching.
#[derive(Debug, Default)]
struct ToolCallBucket {
    id: String,
    name: String,
    arguments: String,
}

fn send_chunk(cx: &ConnectionTo<Client>, session_id: &SessionId, update: SessionUpdate) {
    let notif = SessionNotification::new(session_id.clone(), update);
    if let Err(e) = cx.send_notification(notif) {
        tracing::warn!(error = %format!("{e:#}"), "failed to forward session update");
    }
}

fn text_chunk(text: String) -> agent_client_protocol::schema::ContentChunk {
    use agent_client_protocol::schema::ContentChunk;
    ContentChunk::new(ContentBlock::Text(TextContent::new(text)))
}

/// Mint a synthetic tool_call_id for a malformed `<tool_call>` block.
/// The format mirrors successful calls (`call_<n>`) but uses its own
/// counter so the ids don't collide.
fn next_synthetic_id(counter: &mut usize) -> String {
    let id = format!("call_malformed_{}", *counter);
    *counter += 1;
    id
}

/// Emit a `SessionUpdate::ToolCall` with `Failed` status so Zed
/// renders the malformed call as a structured failure card (raw
/// body visible inside the card) instead of leaving it as inline
/// text in the message pane.
fn emit_malformed_tool_card(
    cx: &ConnectionTo<Client>,
    session_id: &SessionId,
    tool_call_id: &str,
    raw: &str,
) {
    use agent_client_protocol::schema::{
        Content, ToolCall as AcpToolCall, ToolCallContent, ToolCallId, ToolCallStatus, ToolKind,
    };
    let body = format!(
        "Tool call JSON could not be parsed. Raw body:\n\n```\n{raw}\n```\n\n\
         Expected schema:\n\n```json\n{{\"name\": \"<function>\", \"arguments\": {{...}}}}\n```",
    );
    let card = AcpToolCall::new(ToolCallId::new(tool_call_id), "Malformed tool call")
        .kind(ToolKind::Other)
        .status(ToolCallStatus::Failed)
        .raw_input(serde_json::Value::String(raw.to_string()))
        .content(vec![ToolCallContent::Content(Content::new(
            ContentBlock::Text(TextContent::new(body)),
        ))]);
    send_chunk(cx, session_id, SessionUpdate::ToolCall(card));
}

/// Build the assistant-turn / tool-result pair for a malformed
/// `<tool_call>`. The assistant turn carries the raw body verbatim
/// (so the model sees its own previous output), and the tool
/// result spells out *why* it failed with the expected schema —
/// enough for a competent model to self-correct on the next round.
/// Last-chance repair for a malformed `<tool_call>` body: if the
/// model emitted a structurally-valid JSON object with `arguments`
/// but a missing `name`, infer the intended tool from the
/// arguments' shape (see [`tools::infer_tool_name`]). Returns
/// `Some((name, arguments_json))` only when the inference is
/// unambiguous; ambiguous or unrecognised shapes return `None`
/// so the caller surfaces a Failed card.
///
/// We don't try to repair anything qwen3.rs already gave up on for
/// structural reasons (truncation, free-form prose) — those stay
/// Failed and the model retries.
fn try_repair_missing_name(raw: &str) -> Option<(String, String)> {
    let value: serde_json::Value = serde_json::from_str(raw.trim()).ok()?;
    // If a `name` exists at the top level, the parser's own
    // earlier repair passes already had a shot at this and decided
    // it was malformed for some other reason. Don't second-guess
    // them here.
    if value.get("name").is_some() {
        return None;
    }
    let arguments = value.get("arguments")?;
    let name = tools::infer_tool_name(arguments)?;
    let args_json = serde_json::to_string(arguments).ok()?;
    Some((name.to_string(), args_json))
}

fn synthesize_malformed_history(tool_call_id: &str, raw: &str) -> (Message, Message) {
    let call = Message {
        role: Role::Assistant,
        content: MessageContent::ToolCalls {
            text: None,
            calls: vec![ToolCall {
                id: tool_call_id.to_string(),
                // Real tool names never start with `<` — using this
                // placeholder makes the malformed call's identity
                // unambiguous in the rendered transcript.
                name: "<invalid>".to_string(),
                arguments: raw.to_string(),
            }],
        },
    };
    let result = Message {
        role: Role::Tool,
        content: MessageContent::ToolResult {
            tool_call_id: tool_call_id.to_string(),
            content: format!(
                "ERROR: previous <tool_call> body was not valid JSON. Body was:\n{raw}\n\n\
                 Retry with the schema: {{\"name\": \"<function>\", \"arguments\": {{…}}}}"
            ),
        },
    };
    (call, result)
}

/// Compute the prompt token budget for an endpoint given its
/// `context_window` and `max_tokens` settings. The model needs room
/// for both the prompt and its response inside the context window,
/// so the prompt budget is the remainder after subtracting the
/// response cap (defaulting to a conservative 2048 when the endpoint
/// didn't set one) and a small safety margin for tokenizer
/// disagreement.
///
/// The safety margin matters because our per-character estimate in
/// [`compaction`] can drift a few percent from any given upstream
/// tokenizer; we'd rather under-fill the context window than have a
/// well-compacted history still trip `prompt_too_long`.
fn prompt_budget(context_window: usize, max_tokens: Option<u64>) -> usize {
    const SAFETY_MARGIN: usize = 512;
    let max_tokens = max_tokens.unwrap_or(2048) as usize;
    context_window
        .saturating_sub(max_tokens)
        .saturating_sub(SAFETY_MARGIN)
}

fn map_finish_reason(reason: Option<&str>) -> StopReason {
    match reason {
        Some("length") => StopReason::MaxTokens,
        Some("refusal") => StopReason::Refusal,
        // "stop", "tool_calls" (no tools in Stage 2 — degrade to
        // EndTurn so we don't surface a bogus reason), missing, or
        // anything else → EndTurn.
        _ => StopReason::EndTurn,
    }
}

/// Pure helper — turn a prompt's ContentBlocks into the user-message
/// content that goes into history.
///
/// - All-text prompts collapse to [`MessageContent::Text`] (cheaper
///   to encode upstream — many OpenAI-compatible servers prefer the
///   string form when there's no reason to use the array form).
/// - Anything with at least one image becomes
///   [`MessageContent::MultiPart`], preserving block order so the
///   user's "this image, then this text" pacing reaches the model.
/// - `ResourceLink` is rendered as inline text so the model knows
///   it was referenced. Audio and embedded resources aren't
///   advertised as supported in [`PromptCapabilities`]; drop with a
///   warning if a non-conformant client sends one.
fn flatten_prompt(blocks: &[ContentBlock]) -> MessageContent {
    use crate::provider::{ImageData, MessagePart};

    let mut parts: Vec<MessagePart> = Vec::new();
    let mut text_buf = String::new();
    let flush_text = |buf: &mut String, parts: &mut Vec<MessagePart>| {
        if !buf.is_empty() {
            parts.push(MessagePart::Text {
                text: std::mem::take(buf),
            });
        }
    };

    for block in blocks {
        match block {
            ContentBlock::Text(t) => {
                if !text_buf.is_empty() {
                    text_buf.push_str("\n\n");
                }
                text_buf.push_str(&t.text);
            }
            ContentBlock::ResourceLink(link) => {
                if !text_buf.is_empty() {
                    text_buf.push_str("\n\n");
                }
                text_buf.push_str(&format!("[resource link: {}]", link.uri));
            }
            ContentBlock::Image(img) => {
                flush_text(&mut text_buf, &mut parts);
                parts.push(MessagePart::Image(ImageData {
                    mime_type: img.mime_type.clone(),
                    data: img.data.clone(),
                    uri: img.uri.clone(),
                }));
            }
            other => {
                tracing::warn!(?other, "ignoring unsupported content block");
            }
        }
    }
    flush_text(&mut text_buf, &mut parts);

    // Collapse to plain Text when there's no image part — the
    // OpenAI string-form is friendlier to non-vision endpoints
    // (some treat the array form as a vision-only path).
    let has_image = parts.iter().any(|p| matches!(p, MessagePart::Image(_)));
    if !has_image {
        let text = parts
            .into_iter()
            .filter_map(|p| match p {
                MessagePart::Text { text } => Some(text),
                MessagePart::Image(_) => None,
            })
            .collect::<Vec<_>>()
            .join("\n\n");
        return MessageContent::Text { text };
    }
    MessageContent::MultiPart { parts }
}

/// Pure helper — pick which provider handles a session's `model_id`.
/// Returns the matching provider plus the endpoint-local model id
/// (i.e. with any `endpoint:` prefix stripped).
fn resolve_provider(
    providers: &[Arc<dyn Provider>],
    default_endpoint: &str,
    model_id: &str,
) -> anyhow::Result<(Arc<dyn Provider>, String)> {
    let (endpoint_hint, local_model) = parse_model_selector(model_id);
    let target_endpoint = endpoint_hint.unwrap_or(default_endpoint);
    let provider = providers
        .iter()
        .find(|p| p.name() == target_endpoint)
        .ok_or_else(|| anyhow::anyhow!("no provider for endpoint '{target_endpoint}'"))?;
    Ok((provider.clone(), local_model.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use agent_client_protocol::schema::ResourceLink;
    use async_trait::async_trait;
    use futures::stream::BoxStream;

    // ── flatten_prompt ──────────────────────────────────────────────

    fn expect_text(content: &MessageContent) -> &str {
        match content {
            MessageContent::Text { text } => text.as_str(),
            other => panic!("expected MessageContent::Text, got {other:?}"),
        }
    }

    #[test]
    fn flatten_empty_prompt_is_empty() {
        assert_eq!(expect_text(&flatten_prompt(&[])), "");
    }

    #[test]
    fn flatten_joins_text_blocks_with_blank_line() {
        let blocks = vec![
            ContentBlock::Text(TextContent::new("first")),
            ContentBlock::Text(TextContent::new("second")),
        ];
        assert_eq!(expect_text(&flatten_prompt(&blocks)), "first\n\nsecond");
    }

    #[test]
    fn flatten_resource_link_becomes_reference_line() {
        let blocks = vec![ContentBlock::ResourceLink(ResourceLink::new(
            "readme",
            "file:///tmp/x",
        ))];
        assert_eq!(
            expect_text(&flatten_prompt(&blocks)),
            "[resource link: file:///tmp/x]"
        );
    }

    #[test]
    fn flatten_text_and_image_produces_multipart_in_order() {
        use crate::provider::MessagePart;
        let blocks = vec![
            ContentBlock::Text(TextContent::new("describe:")),
            ContentBlock::Image(agent_client_protocol::schema::ImageContent::new(
                "iVBORw0KGgo=",
                "image/png",
            )),
            ContentBlock::Text(TextContent::new("…in detail.")),
        ];
        let content = flatten_prompt(&blocks);
        match content {
            MessageContent::MultiPart { parts } => {
                assert_eq!(parts.len(), 3);
                assert!(matches!(&parts[0], MessagePart::Text { text } if text == "describe:"));
                assert!(matches!(&parts[1], MessagePart::Image(img)
                    if img.mime_type == "image/png" && img.data == "iVBORw0KGgo="));
                assert!(matches!(&parts[2], MessagePart::Text { text } if text == "…in detail."));
            }
            other => panic!("expected MultiPart, got {other:?}"),
        }
    }

    #[test]
    fn flatten_image_only_still_produces_multipart() {
        use crate::provider::MessagePart;
        let blocks = vec![ContentBlock::Image(
            agent_client_protocol::schema::ImageContent::new("AAAA", "image/jpeg"),
        )];
        match flatten_prompt(&blocks) {
            MessageContent::MultiPart { parts } => {
                assert_eq!(parts.len(), 1);
                assert!(matches!(&parts[0], MessagePart::Image(img)
                    if img.mime_type == "image/jpeg"));
            }
            other => panic!("expected MultiPart, got {other:?}"),
        }
    }

    // ── resolve_provider ────────────────────────────────────────────

    /// Minimal Provider stub; just records its name. The trait methods
    /// aren't exercised by resolve_provider so we leave them
    /// unimplemented.
    struct StubProvider(&'static str);

    #[async_trait]
    impl Provider for StubProvider {
        fn name(&self) -> &str {
            self.0
        }
        async fn list_models(&self) -> anyhow::Result<Vec<crate::provider::ModelInfo>> {
            unimplemented!()
        }
        async fn complete(
            &self,
            _request: CompletionRequest,
            _cancel: CancellationToken,
        ) -> anyhow::Result<BoxStream<'static, anyhow::Result<CompletionEvent>>> {
            unimplemented!()
        }
    }

    /// Provider stub whose `list_models` returns canned results.
    /// Used by the `aggregate_models` tests.
    struct ModelProvider {
        name: &'static str,
        models: anyhow::Result<Vec<crate::provider::ModelInfo>>,
    }

    impl ModelProvider {
        fn ok(name: &'static str, ids: &[&str]) -> Arc<dyn Provider> {
            let models = ids
                .iter()
                .map(|id| crate::provider::ModelInfo {
                    id: (*id).to_string(),
                    display_name: None,
                })
                .collect();
            Arc::new(Self {
                name,
                models: Ok(models),
            })
        }
        fn err(name: &'static str, msg: &'static str) -> Arc<dyn Provider> {
            Arc::new(Self {
                name,
                models: Err(anyhow::anyhow!(msg)),
            })
        }
    }

    #[async_trait]
    impl Provider for ModelProvider {
        fn name(&self) -> &str {
            self.name
        }
        async fn list_models(&self) -> anyhow::Result<Vec<crate::provider::ModelInfo>> {
            match &self.models {
                Ok(v) => Ok(v.clone()),
                Err(e) => Err(anyhow::anyhow!("{e:#}")),
            }
        }
        async fn complete(
            &self,
            _request: CompletionRequest,
            _cancel: CancellationToken,
        ) -> anyhow::Result<BoxStream<'static, anyhow::Result<CompletionEvent>>> {
            unimplemented!()
        }
    }

    fn providers() -> Vec<Arc<dyn Provider>> {
        vec![
            Arc::new(StubProvider("helexa")),
            Arc::new(StubProvider("openrouter")),
        ]
    }

    #[test]
    fn bare_model_routes_to_default() {
        let (p, m) = resolve_provider(&providers(), "helexa", "helexa/large").unwrap();
        assert_eq!(p.name(), "helexa");
        assert_eq!(m, "helexa/large");
    }

    #[test]
    fn prefixed_model_routes_by_endpoint() {
        let (p, m) =
            resolve_provider(&providers(), "helexa", "openrouter:anthropic/claude-opus-4").unwrap();
        assert_eq!(p.name(), "openrouter");
        assert_eq!(m, "anthropic/claude-opus-4");
    }

    #[test]
    fn unknown_endpoint_errors() {
        // `Arc<dyn Provider>` doesn't impl Debug, which rules out
        // `.unwrap_err()` (it requires T: Debug). Pattern-match instead.
        match resolve_provider(&providers(), "helexa", "ghost:gpt-9") {
            Ok(_) => panic!("expected error for unknown endpoint"),
            Err(e) => assert!(format!("{e}").contains("ghost")),
        }
    }

    // ── map_finish_reason ───────────────────────────────────────────

    // ── prompt_budget ───────────────────────────────────────────────

    #[test]
    fn prompt_budget_reserves_response_and_safety() {
        // 32K window, 8K response cap → 32768 - 8192 - 512 = 24064.
        assert_eq!(prompt_budget(32_768, Some(8_192)), 24_064);
    }

    #[test]
    fn prompt_budget_uses_default_when_max_tokens_unset() {
        // Default response cap = 2048; safety = 512.
        assert_eq!(prompt_budget(32_768, None), 32_768 - 2_048 - 512);
    }

    #[test]
    fn prompt_budget_saturates_when_window_too_small() {
        // Pathological config: window smaller than response + safety.
        // Don't underflow — return zero so compaction tries hardest
        // and upstream surfaces the inevitable error.
        assert_eq!(prompt_budget(1_000, Some(8_192)), 0);
    }

    // ── aggregate_models ────────────────────────────────────────────

    #[tokio::test]
    async fn aggregate_models_single_endpoint_has_bare_ids() {
        let providers = vec![ModelProvider::ok(
            "helexa",
            &["helexa/large", "helexa/small"],
        )];
        let models = aggregate_models(&providers).await;
        let ids: Vec<&str> = models.iter().map(|m| m.model_id.0.as_ref()).collect();
        assert_eq!(ids, vec!["helexa/large", "helexa/small"]);
    }

    #[tokio::test]
    async fn aggregate_models_multi_endpoint_prefixes_every_id() {
        let providers = vec![
            ModelProvider::ok("helexa", &["helexa/large"]),
            ModelProvider::ok("openrouter", &["anthropic/claude-opus-4"]),
        ];
        let models = aggregate_models(&providers).await;
        let ids: Vec<&str> = models.iter().map(|m| m.model_id.0.as_ref()).collect();
        assert_eq!(
            ids,
            vec!["helexa:helexa/large", "openrouter:anthropic/claude-opus-4"]
        );
    }

    #[tokio::test]
    async fn aggregate_models_skips_failing_endpoint() {
        let providers = vec![
            ModelProvider::err("flaky", "boom"),
            ModelProvider::ok("openrouter", &["gpt-9"]),
        ];
        let models = aggregate_models(&providers).await;
        let ids: Vec<&str> = models.iter().map(|m| m.model_id.0.as_ref()).collect();
        // Multi-endpoint case → prefix survives even when one
        // endpoint dropped out. flaky's models are absent, not
        // null-filled.
        assert_eq!(ids, vec!["openrouter:gpt-9"]);
    }

    #[test]
    fn maps_known_finish_reasons() {
        assert!(matches!(
            map_finish_reason(Some("length")),
            StopReason::MaxTokens
        ));
        assert!(matches!(
            map_finish_reason(Some("refusal")),
            StopReason::Refusal
        ));
        assert!(matches!(
            map_finish_reason(Some("stop")),
            StopReason::EndTurn
        ));
        assert!(matches!(
            map_finish_reason(Some("tool_calls")),
            StopReason::EndTurn
        ));
        assert!(matches!(map_finish_reason(None), StopReason::EndTurn));
    }
}
