// LlamaCppClient — talks to a llama-swap proxy that fronts one or more
// llama-server processes. llama-swap exposes an OpenAI-compatible HTTP
// surface (`/v1/chat/completions`, `/v1/embeddings`, `/v1/models`), so the
// wire translation mirrors `OpenRouterClient` almost exactly.
//
// Differences from OpenRouter:
// - No bearer auth or attribution headers; llama-swap is LAN-only.
// - Three model slots (`primary_model` = chat, `vision_model`, `embedding_model`)
//   each map to a model id in the llama-swap config. `describe_image` and
//   `generate_embeddings` issue requests with the appropriate slot id in the
//   `model` field, which is how llama-swap selects which backend process to
//   run.
// - `/v1/models` returns only the configured slot ids — capabilities aren't
//   reported by the API, so we infer `has_vision` from a single config field
//   (`vision_model`, defaulting to `"vision"`) and assume `has_tool_calling`
//   is true for every slot, since llama-swap entries default to launching
//   llama-server with `--jinja`.
//
// First consumer lands alongside the three-way backend dispatch in
// insight_generator / insight_chat.
#![allow(dead_code)]

use anyhow::Context as _;
use anyhow::{Result, anyhow, bail};
use async_trait::async_trait;
use opentelemetry::trace::Status;
use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json};
use std::collections::HashMap;
use std::time::Duration;
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

use crate::ai::llm_client::{
    ChatMessage, LlmClient, LlmStreamEvent, ModelCapabilities, Tool, ToolCall, ToolCallFunction,
};
use futures::stream::{BoxStream, StreamExt};

const DEFAULT_BASE_URL: &str = "http://localhost:9292/v1";
const DEFAULT_PRIMARY_MODEL: &str = "chat";
const DEFAULT_VISION_MODEL: &str = "vision";
const DEFAULT_EMBEDDING_MODEL: &str = "embed";
const DEFAULT_TTS_MODEL: &str = "chatterbox";
const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 180;
/// How long a served-context-window lookup stays cached per model.
const CONTEXT_WINDOW_CACHE_TTL: Duration = Duration::from_secs(600);

/// `"<swap root>|<model>"` -> (n_ctx, fetched at). Process-wide because
/// clients are cloned per request.
fn context_window_cache() -> &'static std::sync::Mutex<HashMap<String, (usize, std::time::Instant)>>
{
    static CACHE: std::sync::OnceLock<
        std::sync::Mutex<HashMap<String, (usize, std::time::Instant)>>,
    > = std::sync::OnceLock::new();
    CACHE.get_or_init(Default::default)
}

/// OpenAI-compatible client targeting a llama-swap proxy in front of one or
/// more llama-server processes. See the module doc-comment for the slot model.
#[derive(Clone)]
pub struct LlamaCppClient {
    client: Client,
    pub base_url: String,
    /// Chat model slot id (e.g. `"chat"`). Used for `generate` /
    /// `chat_with_tools` / `chat_with_tools_stream`.
    pub primary_model: String,
    /// Embedding model slot id (e.g. `"embed"`). Used for
    /// `generate_embeddings`.
    pub embedding_model: String,
    /// Vision model slot id. Used for `describe_image` routing. Defaults
    /// to `primary_model` so describe_image works out of the box; override
    /// via `LLAMA_SWAP_VISION_MODEL` for a dedicated vision slot.
    pub vision_model: String,
    /// TTS model slot id (e.g. `"chatterbox"`). Routes `text_to_speech` and
    /// is the `/upstream/<id>/voices` path segment for the voice library.
    /// Override via `LLAMA_SWAP_TTS_MODEL`.
    pub tts_model: String,
    num_ctx: Option<i32>,
    temperature: Option<f32>,
    top_p: Option<f32>,
    top_k: Option<i32>,
    min_p: Option<f32>,
    /// When `Some`, forwarded to llama-server as
    /// `chat_template_kwargs: {"enable_thinking": <bool>}`. The Jinja chat
    /// template (e.g. Qwen3) reads this to gate its reasoning block. `None`
    /// omits the key entirely, leaving the template's own default. Templates
    /// that don't reference the key ignore it, so sending it is harmless.
    enable_thinking: Option<bool>,
}

impl LlamaCppClient {
    pub fn new(base_url: Option<String>, primary_model: Option<String>) -> Self {
        let timeout_secs = std::env::var("LLAMA_SWAP_REQUEST_TIMEOUT_SECONDS")
            .ok()
            .and_then(|v| v.parse::<u64>().ok())
            .unwrap_or(DEFAULT_REQUEST_TIMEOUT_SECS);
        let pm = primary_model.unwrap_or_else(|| DEFAULT_PRIMARY_MODEL.to_string());
        Self {
            client: Client::builder()
                .connect_timeout(Duration::from_secs(10))
                .timeout(Duration::from_secs(timeout_secs))
                .build()
                .unwrap_or_else(|_| Client::new()),
            base_url: base_url.unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
            primary_model: pm.clone(),
            embedding_model: DEFAULT_EMBEDDING_MODEL.to_string(),
            vision_model: pm,
            tts_model: DEFAULT_TTS_MODEL.to_string(),
            num_ctx: None,
            temperature: None,
            top_p: None,
            top_k: None,
            min_p: None,
            enable_thinking: None,
        }
    }

    pub fn set_embedding_model(&mut self, model: String) {
        self.embedding_model = model;
    }

    pub fn set_vision_model(&mut self, model: String) {
        self.vision_model = model;
    }

    pub fn set_num_ctx(&mut self, num_ctx: Option<i32>) {
        self.num_ctx = num_ctx;
    }

    /// Set the reasoning toggle forwarded as `chat_template_kwargs.enable_thinking`.
    /// `None` leaves the chat template's own default in place.
    pub fn set_enable_thinking(&mut self, enable_thinking: Option<bool>) {
        self.enable_thinking = enable_thinking;
    }

    pub fn set_sampling_params(
        &mut self,
        temperature: Option<f32>,
        top_p: Option<f32>,
        top_k: Option<i32>,
        min_p: Option<f32>,
    ) {
        self.temperature = temperature;
        self.top_p = top_p;
        self.top_k = top_k;
        self.min_p = min_p;
    }

    pub fn set_tts_model(&mut self, model: String) {
        self.tts_model = model;
    }

    // --- TTS (Chatterbox behind llama-swap) ---------------------------------
    //
    // Speech synthesis uses the OpenAI-compatible `{base_url}/audio/speech`
    // endpoint (llama-swap routes by the `model` field). The voice *library*
    // (list / create cloned voices) is NOT an OpenAI endpoint — it lives on the
    // upstream server directly, reached via llama-swap's passthrough at
    // `{swap_root}/upstream/<tts_model>/voices`.

    /// Per-request context window `model` is actually served with, read from
    /// llama-server's `/props` (`default_generation_settings.n_ctx`, the
    /// per-slot window) through llama-swap's passthrough. That request loads
    /// the model if it isn't already, which is fine: callers ask right before
    /// chatting with it. Successful lookups are cached for
    /// `CONTEXT_WINDOW_CACHE_TTL` so a config.yaml edit is picked up without
    /// a restart; failures aren't cached and return `None`.
    pub async fn served_context_window(&self, model: &str) -> Option<usize> {
        if let Some(n_ctx) = self.cached_context_window(model) {
            return Some(n_ctx);
        }
        let key = format!("{}|{}", self.swap_root(), model);
        let url = format!("{}/upstream/{}/props", self.swap_root(), model);
        let props: Value = match self.client.get(&url).send().await {
            Ok(resp) if resp.status().is_success() => resp.json().await.ok()?,
            Ok(resp) => {
                tracing::warn!(model, status = %resp.status(), "llama-swap props lookup failed");
                return None;
            }
            Err(e) => {
                tracing::warn!(model, error = %e, "llama-swap props lookup failed");
                return None;
            }
        };
        let n_ctx = props
            .pointer("/default_generation_settings/n_ctx")
            .or_else(|| props.get("n_ctx"))
            .and_then(Value::as_u64)
            .filter(|n| *n > 0)? as usize;
        if let Ok(mut cache) = context_window_cache().lock() {
            cache.insert(key, (n_ctx, std::time::Instant::now()));
        }
        Some(n_ctx)
    }

    /// Served context windows for `models`, for display, WITHOUT loading
    /// anything: a model counts only if its window is cached from an
    /// earlier lookup or llama-swap reports it as currently loaded (then
    /// `/props` is answered by the running process). Others are absent.
    pub async fn known_context_windows(&self, models: &[String]) -> HashMap<String, usize> {
        let loaded: std::collections::HashSet<String> = async {
            let url = format!("{}/models", self.base_url);
            let parsed: Value = self.client.get(&url).send().await.ok()?.json().await.ok()?;
            Some(
                parsed
                    .get("data")?
                    .as_array()?
                    .iter()
                    .filter(|m| {
                        m.pointer("/status/value").and_then(Value::as_str) == Some("loaded")
                    })
                    .filter_map(|m| m.get("id")?.as_str().map(str::to_string))
                    .collect(),
            )
        }
        .await
        .unwrap_or_default();

        let mut out = HashMap::new();
        for model in models {
            let cached = self.cached_context_window(model);
            let n_ctx = match cached {
                Some(n) => Some(n),
                None if loaded.contains(model) => self.served_context_window(model).await,
                None => None,
            };
            if let Some(n) = n_ctx {
                out.insert(model.clone(), n);
            }
        }
        out
    }

    fn cached_context_window(&self, model: &str) -> Option<usize> {
        let key = format!("{}|{}", self.swap_root(), model);
        let cache = context_window_cache().lock().ok()?;
        let &(n_ctx, at) = cache.get(&key)?;
        (at.elapsed() < CONTEXT_WINDOW_CACHE_TTL).then_some(n_ctx)
    }

    /// Root of the llama-swap proxy: `base_url` with a trailing `/v1` removed.
    /// The `/upstream/...` passthrough lives here, not under `/v1`.
    fn swap_root(&self) -> &str {
        let b = self.base_url.trim_end_matches('/');
        b.strip_suffix("/v1").unwrap_or(b)
    }

    /// Synthesize speech for `input` in an optional named `voice`, returning
    /// the raw audio bytes (format per `response_format`, e.g. `"mp3"`/`"wav"`).
    ///
    /// Chatterbox generation knobs are forwarded when set (caller is expected
    /// to have range-clamped them): `exaggeration` (0.25–2.0, emotion),
    /// `cfg_weight` (0.0–1.0, pace), `temperature` (0.05–5.0, randomness).
    ///
    /// Callers must hold the GPU write lease (`ai::gpu::tts_lease`) across
    /// this call. It is taken at the call sites in `ai::tts` rather than here
    /// so the speech-job path can flip its job to `running` between acquiring
    /// the GPU and sending the request.
    pub async fn text_to_speech(
        &self,
        input: &str,
        voice: Option<&str>,
        response_format: &str,
        exaggeration: Option<f32>,
        cfg_weight: Option<f32>,
        temperature: Option<f32>,
    ) -> Result<Vec<u8>> {
        let url = format!("{}/audio/speech", self.base_url);
        let mut body = json!({
            "model": self.tts_model,
            "input": input,
            "response_format": response_format,
        });
        if let Some(v) = voice {
            body["voice"] = Value::String(v.to_string());
        }
        if let Some(x) = exaggeration {
            body["exaggeration"] = json!(x);
        }
        if let Some(x) = cfg_weight {
            body["cfg_weight"] = json!(x);
        }
        if let Some(x) = temperature {
            body["temperature"] = json!(x);
        }

        // TTS gets its own (longer) timeout: synthesizing a long, internally
        // chunked insight can take minutes, well past the shared chat/embedding
        // client timeout. Per-request `.timeout()` overrides the client default.
        let tts_timeout = std::env::var("LLAMA_SWAP_TTS_REQUEST_TIMEOUT_SECONDS")
            .ok()
            .and_then(|v| v.parse::<u64>().ok())
            .filter(|n| *n > 0)
            .unwrap_or(600);

        let resp = self
            .client
            .post(&url)
            .timeout(Duration::from_secs(tts_timeout))
            .json(&body)
            .send()
            .await
            .with_context(|| format!("POST {} failed", url))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            bail!("llama-swap TTS request failed: {} — {}", status, text);
        }

        Ok(resp
            .bytes()
            .await
            .context("reading TTS audio bytes")?
            .to_vec())
    }

    /// List voices in the Chatterbox voice library (raw JSON passthrough).
    pub async fn list_voices(&self) -> Result<Value> {
        let url = format!("{}/upstream/{}/voices", self.swap_root(), self.tts_model);
        // The /upstream passthrough spins Chatterbox up (evicting the LLM),
        // so it takes the exclusive GPU lease like synthesis does.
        let _gpu = crate::ai::gpu::tts_lease().await;
        let resp = self
            .client
            .get(&url)
            .send()
            .await
            .with_context(|| format!("GET {} failed", url))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            bail!("llama-swap list_voices failed: {} — {}", status, text);
        }
        resp.json().await.context("parsing voices response")
    }

    /// Register a cloned voice from raw audio bytes (multipart `voice_name` +
    /// `voice_file`). Returns the upstream JSON response.
    pub async fn create_voice(
        &self,
        voice_name: &str,
        audio_bytes: Vec<u8>,
        filename: &str,
        mime: &str,
    ) -> Result<Value> {
        let url = format!("{}/upstream/{}/voices", self.swap_root(), self.tts_model);
        let part = reqwest::multipart::Part::bytes(audio_bytes)
            .file_name(filename.to_string())
            .mime_str(mime)
            .context("invalid audio mime type")?;
        let form = reqwest::multipart::Form::new()
            .text("voice_name", voice_name.to_string())
            .part("voice_file", part);

        // The /upstream passthrough spins Chatterbox up (evicting the LLM),
        // so it takes the exclusive GPU lease like synthesis does.
        let _gpu = crate::ai::gpu::tts_lease().await;
        let resp = self
            .client
            .post(&url)
            .multipart(form)
            .send()
            .await
            .with_context(|| format!("POST {} (multipart) failed", url))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            bail!("llama-swap create_voice failed: {} — {}", status, text);
        }
        resp.json().await.context("parsing create_voice response")
    }

    /// Delete a cloned voice from the Chatterbox voice library
    /// (`DELETE /voices/{name}` on the upstream, via llama-swap passthrough).
    pub async fn delete_voice(&self, voice_name: &str) -> Result<Value> {
        let url = format!(
            "{}/upstream/{}/voices/{}",
            self.swap_root(),
            self.tts_model,
            voice_name
        );
        // The /upstream passthrough spins Chatterbox up (evicting the LLM),
        // so it takes the exclusive GPU lease like synthesis does.
        let _gpu = crate::ai::gpu::tts_lease().await;
        let resp = self
            .client
            .delete(&url)
            .send()
            .await
            .with_context(|| format!("DELETE {} failed", url))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            bail!("llama-swap delete_voice failed: {} — {}", status, text);
        }
        // Some upstreams reply with an empty body on delete.
        Ok(resp
            .json()
            .await
            .unwrap_or_else(|_| json!({ "status": "deleted" })))
    }

    /// Translate canonical messages to the OpenAI-compatible wire shape.
    /// Behaviorally identical to `OpenRouterClient::messages_to_openai` —
    /// stringify tool-call arguments, rewrite images into content-parts, attach
    /// `tool_call_id` to `role=tool` messages based on the preceding assistant
    /// turn's tool calls.
    fn messages_to_openai(messages: &[ChatMessage]) -> Vec<Value> {
        let mut out = Vec::with_capacity(messages.len());
        let mut last_tool_call_ids: Vec<String> = Vec::new();
        let mut next_tool_result_idx: usize = 0;

        for msg in messages {
            let mut obj = serde_json::Map::new();
            obj.insert("role".into(), Value::String(msg.role.clone()));

            // Assistant messages with tool_calls must emit `content: null`
            // (not `""`) — some Jinja templates (Mistral-family) treat
            // empty-string content as a regular message rather than a
            // tool-calling turn, breaking role-alternation validation.
            let has_tool_calls = msg.role == "assistant"
                && msg.tool_calls.as_ref().is_some_and(|tcs| !tcs.is_empty());

            match &msg.images {
                Some(images) if !images.is_empty() => {
                    let mut parts: Vec<Value> = Vec::new();
                    if !msg.content.is_empty() {
                        parts.push(json!({"type": "text", "text": msg.content}));
                    }
                    for img in images {
                        let url = image_to_data_url(img);
                        parts.push(json!({
                            "type": "image_url",
                            "image_url": { "url": url }
                        }));
                    }
                    obj.insert("content".into(), Value::Array(parts));
                }
                _ if has_tool_calls && msg.content.is_empty() => {
                    obj.insert("content".into(), Value::Null);
                }
                _ => {
                    obj.insert("content".into(), Value::String(msg.content.clone()));
                }
            }

            if let Some(tcs) = &msg.tool_calls
                && msg.role == "assistant"
            {
                let converted: Vec<Value> = tcs
                    .iter()
                    .enumerate()
                    .map(|(i, call)| {
                        let id = call.id.clone().unwrap_or_else(|| format!("call_{}", i));
                        let args_str = serde_json::to_string(&call.function.arguments)
                            .unwrap_or_else(|_| "{}".to_string());
                        json!({
                            "id": id,
                            "type": "function",
                            "function": {
                                "name": call.function.name,
                                "arguments": args_str,
                            }
                        })
                    })
                    .collect();
                last_tool_call_ids = converted
                    .iter()
                    .filter_map(|v| v.get("id").and_then(|x| x.as_str()).map(String::from))
                    .collect();
                next_tool_result_idx = 0;
                obj.insert("tool_calls".into(), Value::Array(converted));
            }

            if msg.role == "tool" {
                let id = last_tool_call_ids
                    .get(next_tool_result_idx)
                    .cloned()
                    .unwrap_or_else(|| "call_0".to_string());
                obj.insert("tool_call_id".into(), Value::String(id));
                next_tool_result_idx += 1;
            }

            out.push(Value::Object(obj));
        }

        out
    }

    /// Parse an OpenAI-compatible assistant message back into canonical shape.
    /// llama.cpp emits `reasoning_content` on thinking models; we drop it for
    /// parity with OpenRouter (which also strips upstream reasoning fields).
    fn openai_message_to_chat(msg: &Value) -> Result<ChatMessage> {
        let obj = msg
            .as_object()
            .ok_or_else(|| anyhow!("response message is not an object"))?;
        let role = obj
            .get("role")
            .and_then(|v| v.as_str())
            .unwrap_or("assistant")
            .to_string();
        let content = obj
            .get("content")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();

        let tool_calls = if let Some(tcs) = obj.get("tool_calls").and_then(|v| v.as_array()) {
            let mut parsed = Vec::with_capacity(tcs.len());
            for tc in tcs {
                let id = tc.get("id").and_then(|v| v.as_str()).map(String::from);
                let function = tc
                    .get("function")
                    .ok_or_else(|| anyhow!("tool_call missing function field"))?;
                let name = function
                    .get("name")
                    .and_then(|v| v.as_str())
                    .unwrap_or_default()
                    .to_string();
                let args_value = match function.get("arguments") {
                    Some(Value::String(s)) => {
                        serde_json::from_str::<Value>(s).unwrap_or_else(|_| json!({}))
                    }
                    Some(v @ Value::Object(_)) => v.clone(),
                    _ => json!({}),
                };
                parsed.push(ToolCall {
                    id,
                    function: ToolCallFunction {
                        name,
                        arguments: args_value,
                    },
                });
            }
            Some(parsed)
        } else {
            None
        };

        Ok(ChatMessage {
            role,
            content,
            tool_calls,
            images: None,
        })
    }

    fn build_options(&self) -> Vec<(&'static str, Value)> {
        let mut v = Vec::new();
        if let Some(t) = self.temperature {
            v.push(("temperature", json!(t)));
        }
        if let Some(p) = self.top_p {
            v.push(("top_p", json!(p)));
        }
        if let Some(k) = self.top_k {
            v.push(("top_k", json!(k)));
        }
        if let Some(m) = self.min_p {
            v.push(("min_p", json!(m)));
        }
        // num_ctx isn't an OpenAI param; llama-server bakes ctx in at launch
        // via -c, so we silently drop the override here. The config.yaml
        // entry is the source of truth for context size; chat budgets read
        // it back via `context_window()`.
        let _ = self.num_ctx;
        // Reasoning toggle for thinking-capable templates (Qwen3 et al.).
        // llama-server forwards chat_template_kwargs into the Jinja render
        // (requires --jinja); templates that ignore the key are unaffected.
        if let Some(think) = self.enable_thinking {
            v.push(("chat_template_kwargs", json!({ "enable_thinking": think })));
        }
        v
    }

    /// Issue a chat request with an explicit model id override. Used by
    /// `describe_image` to route through the vision slot without mutating
    /// `self.primary_model`.
    async fn chat_completion_with_model(
        &self,
        model: &str,
        span: Option<&tracing::Span>,
        messages: Vec<ChatMessage>,
        tools: Vec<Tool>,
    ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
        let url = format!("{}/chat/completions", self.base_url);
        let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect();
        log::debug!(
            "llama-swap chat_completion: model={} roles={:?} tools={}",
            model,
            roles,
            tools.len()
        );
        let mut body = serde_json::Map::new();
        body.insert("model".into(), Value::String(model.to_string()));
        body.insert(
            "messages".into(),
            Value::Array(Self::messages_to_openai(&messages)),
        );
        body.insert("stream".into(), Value::Bool(false));
        if !tools.is_empty() {
            body.insert(
                "tools".into(),
                serde_json::to_value(&tools).context("serializing tools")?,
            );
        }
        for (k, v) in self.build_options() {
            body.insert(k.into(), v);
        }

        // Wait for any TTS synthesis to release the GPU before the request
        // timeout starts (see ai::gpu).
        let _gpu = crate::ai::gpu::llm_lease().await;
        let mut request = self.client.post(&url).json(&Value::Object(body));
        if let Some(s) = span {
            request = crate::otel::with_traceparent(request, s);
        }
        let resp = match request.send().await {
            Ok(r) => r,
            Err(e) => {
                if let Some(s) = span {
                    s.set_status(Status::error("llama-swap request failed"));
                }
                return Err(e).with_context(|| format!("POST {} failed", url));
            }
        };

        if !resp.status().is_success() {
            let status = resp.status();
            crate::otel::record_http_failure(
                span,
                status.as_u16(),
                &format!("llama-swap chat request failed: {status}"),
            );
            let body = resp.text().await.unwrap_or_default();
            bail!("llama-swap chat request failed: {} — {}", status, body);
        }
        crate::otel::set_http_status(span, resp.status().as_u16());

        let parsed: Value = resp.json().await.context("parsing chat response")?;
        let choice = parsed
            .get("choices")
            .and_then(|v| v.as_array())
            .and_then(|a| a.first())
            .ok_or_else(|| {
                anyhow!(
                    "response missing choices[0]: {}",
                    extract_error_detail(&parsed)
                )
            })?;
        let msg = choice.get("message").ok_or_else(|| {
            anyhow!(
                "choices[0] missing message: {}",
                extract_error_detail(&parsed)
            )
        })?;
        let chat_msg = Self::openai_message_to_chat(msg)?;

        let usage = parsed.get("usage");
        let prompt_tokens = usage
            .and_then(|u| u.get("prompt_tokens"))
            .and_then(|v| v.as_i64())
            .map(|n| n as i32);
        let completion_tokens = usage
            .and_then(|u| u.get("completion_tokens"))
            .and_then(|v| v.as_i64())
            .map(|n| n as i32);

        log_timings(&parsed, prompt_tokens, completion_tokens);

        if let Some(s) = span {
            if let Some(pt) = prompt_tokens {
                s.set_attribute("llm.tokens.input", pt as i64);
            }
            if let Some(ct) = completion_tokens {
                s.set_attribute("llm.tokens.output", ct as i64);
            }
            s.set_status(Status::Ok);
        }

        Ok((chat_msg, prompt_tokens, completion_tokens))
    }

    /// Start the `llm.request` span for an outgoing llama-swap call.
    fn start_request_span(&self, operation: &str, model: &str) -> tracing::Span {
        crate::otel::start_llm_request_span("llamacpp", operation, model)
    }
}

#[async_trait]
impl LlmClient for LlamaCppClient {
    async fn generate(
        &self,
        prompt: &str,
        system: Option<&str>,
        images: Option<Vec<String>>,
    ) -> Result<String> {
        let mut messages: Vec<ChatMessage> = Vec::new();
        if let Some(sys) = system {
            messages.push(ChatMessage::system(sys));
        }
        let mut user = ChatMessage::user(prompt);
        user.images = images;
        messages.push(user);

        // Delegates to chat_with_tools, which starts the llm.request span
        // (operation=chat) — one span per HTTP request.
        let (reply, _, _) = self.chat_with_tools(messages, Vec::new()).await?;
        Ok(reply.content)
    }

    async fn chat_with_tools(
        &self,
        messages: Vec<ChatMessage>,
        tools: Vec<Tool>,
    ) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
        log::info!(
            "llama-swap chat_with_tools: model={} messages={} tools={}",
            self.primary_model,
            messages.len(),
            tools.len()
        );
        let span = self.start_request_span("chat", &self.primary_model);
        self.chat_completion_with_model(&self.primary_model.clone(), Some(&span), messages, tools)
            .await
    }

    async fn chat_with_tools_stream(
        &self,
        messages: Vec<ChatMessage>,
        tools: Vec<Tool>,
    ) -> Result<BoxStream<'static, Result<LlmStreamEvent>>> {
        // The span outlives this method — it moves into the stream
        // generator and is finalized when the response drains.
        let span = self.start_request_span("chat_stream", &self.primary_model);
        let url = format!("{}/chat/completions", self.base_url);
        let roles: Vec<&str> = messages.iter().map(|m| m.role.as_str()).collect();
        log::debug!(
            "llama-swap stream: model={} roles={:?} tools={}",
            self.primary_model,
            roles,
            tools.len()
        );
        let mut body = serde_json::Map::new();
        body.insert("model".into(), Value::String(self.primary_model.clone()));
        body.insert(
            "messages".into(),
            Value::Array(Self::messages_to_openai(&messages)),
        );
        body.insert("stream".into(), Value::Bool(true));
        body.insert(
            "stream_options".into(),
            serde_json::json!({ "include_usage": true }),
        );
        if !tools.is_empty() {
            body.insert(
                "tools".into(),
                serde_json::to_value(&tools).context("serializing tools")?,
            );
        }
        for (k, v) in self.build_options() {
            body.insert(k.into(), v);
        }

        // Wait for any TTS synthesis to release the GPU before the request
        // timeout starts (see ai::gpu). The guard is moved into the stream
        // below so the lease spans the whole generation, not just the send.
        let gpu = crate::ai::gpu::llm_lease().await;
        let mut request = self.client.post(&url).json(&Value::Object(body));
        request = crate::otel::with_traceparent(request, &span);
        let resp = request
            .send()
            .await
            .with_context(|| format!("POST {} failed", url))?;

        if !resp.status().is_success() {
            let status = resp.status();
            crate::otel::record_http_failure(
                Some(&span),
                status.as_u16(),
                &format!("llama-swap stream request failed: {status}"),
            );
            let body = resp.text().await.unwrap_or_default();
            bail!("llama-swap stream request failed: {} — {}", status, body);
        }
        crate::otel::set_http_status(Some(&span), resp.status().as_u16());

        let byte_stream = resp.bytes_stream();
        let stream = async_stream::stream! {
            let _gpu = gpu;
            let mut byte_stream = byte_stream;
            let mut buf: Vec<u8> = Vec::new();
            let mut accumulated_content = String::new();
            let mut tool_state = ToolCallAssembly::new();
            let mut role = "assistant".to_string();
            let mut prompt_tokens: Option<i32> = None;
            let mut completion_tokens: Option<i32> = None;
            let mut last_frame: Option<Value> = None;
            let mut done_seen = false;

            while let Some(chunk) = byte_stream.next().await {
                let chunk = match chunk {
                    Ok(b) => b,
                    Err(e) => {
                        yield Err(anyhow!("stream read failed: {}", e));
                        return;
                    }
                };
                buf.extend_from_slice(&chunk);

                while let Some(sep) = find_double_newline(&buf) {
                    let frame = buf.drain(..sep + 2).collect::<Vec<_>>();
                    let frame_str = match std::str::from_utf8(&frame) {
                        Ok(s) => s,
                        Err(_) => continue,
                    };
                    for line in frame_str.lines() {
                        let line = line.trim_end_matches('\r');
                        let payload = match line.strip_prefix("data: ") {
                            Some(p) => p,
                            None => continue,
                        };
                        if payload == "[DONE]" {
                            done_seen = true;
                            break;
                        }
                        let v: Value = match serde_json::from_str(payload) {
                            Ok(v) => v,
                            Err(e) => {
                                log::warn!(
                                    "malformed llama-swap SSE frame: {} ({})",
                                    payload,
                                    e
                                );
                                continue;
                            }
                        };

                        if let Some(usage) = v.get("usage") {
                            prompt_tokens = usage
                                .get("prompt_tokens")
                                .and_then(|n| n.as_i64())
                                .map(|n| n as i32);
                            completion_tokens = usage
                                .get("completion_tokens")
                                .and_then(|n| n.as_i64())
                                .map(|n| n as i32);
                            last_frame = Some(v.clone());
                        }

                        let Some(choices) = v.get("choices").and_then(|c| c.as_array())
                        else {
                            continue;
                        };
                        let Some(choice) = choices.first() else { continue };
                        let delta = match choice.get("delta") {
                            Some(d) => d,
                            None => continue,
                        };
                        if let Some(r) = delta.get("role").and_then(|v| v.as_str()) {
                            role = r.to_string();
                        }
                        if let Some(content) =
                            delta.get("content").and_then(|v| v.as_str())
                            && !content.is_empty()
                        {
                            accumulated_content.push_str(content);
                            yield Ok(LlmStreamEvent::TextDelta(content.to_string()));
                        }
                        if let Some(tcs) = delta.get("tool_calls").and_then(|v| v.as_array()) {
                            apply_tool_call_deltas(&mut tool_state, tcs);
                        }
                    }
                    if done_seen {
                        break;
                    }
                }
                if done_seen {
                    break;
                }
            }

            let tool_calls = finalize_tool_calls(tool_state);

            if let Some(ref frame) = last_frame {
                log_timings(frame, prompt_tokens, completion_tokens);
            }

            let message = ChatMessage {
                role,
                content: accumulated_content,
                tool_calls,
                images: None,
            };
            if let Some(pt) = prompt_tokens {
                span.set_attribute("llm.tokens.input", pt as i64);
            }
            if let Some(ct) = completion_tokens {
                span.set_attribute("llm.tokens.output", ct as i64);
            }
            span.set_status(Status::Ok);
            yield Ok(LlmStreamEvent::Done {
                message,
                prompt_eval_count: prompt_tokens,
                eval_count: completion_tokens,
            });
        };

        Ok(Box::pin(stream))
    }

    async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        let span = self.start_request_span("embeddings", &self.embedding_model);
        // Deliberately NO GPU lease: the embed slot sits in llama-swap's
        // always-resident group and never participates in a model swap, so
        // leasing here would only stall searches behind a queued synthesis.
        let url = format!("{}/embeddings", self.base_url);
        let body = json!({
            "model": self.embedding_model,
            "input": texts,
        });

        let mut request = self.client.post(&url).json(&body);
        request = crate::otel::with_traceparent(request, &span);
        let resp = match request.send().await {
            Ok(r) => r,
            Err(e) => {
                span.set_status(Status::error("llama-swap embedding request failed"));
                return Err(e).with_context(|| format!("POST {} failed", url));
            }
        };

        if !resp.status().is_success() {
            let status = resp.status();
            crate::otel::record_http_failure(
                Some(&span),
                status.as_u16(),
                &format!("llama-swap embedding request failed: {status}"),
            );
            let body = resp.text().await.unwrap_or_default();
            bail!("llama-swap embedding request failed: {} — {}", status, body);
        }
        crate::otel::set_http_status(Some(&span), resp.status().as_u16());

        #[derive(Deserialize)]
        struct EmbedResponse {
            data: Vec<EmbedItem>,
        }
        #[derive(Deserialize)]
        struct EmbedItem {
            embedding: Vec<f32>,
        }

        let parsed: EmbedResponse = match resp.json().await {
            Ok(p) => p,
            Err(e) => {
                span.set_status(Status::error("parsing embed response failed"));
                return Err(e).context("parsing embed response");
            }
        };
        span.set_status(Status::Ok);
        Ok(parsed.data.into_iter().map(|i| i.embedding).collect())
    }

    async fn describe_image(&self, image_base64: &str) -> Result<String> {
        // Own span with operation=describe_image — routes through the vision slot.
        let span = self.start_request_span("describe_image", &self.vision_model);
        let prompt = "Briefly describe what you see in this image in 1-2 sentences. \
                      Focus on the people, location, and activity.";
        let system = "You are a scene description assistant. Be concise and factual.";

        let messages = vec![
            ChatMessage::system(system),
            ChatMessage {
                role: "user".to_string(),
                content: prompt.to_string(),
                tool_calls: None,
                images: Some(vec![image_base64.to_string()]),
            },
        ];

        let (reply, _, _) = self
            .chat_completion_with_model(
                &self.vision_model.clone(),
                Some(&span),
                messages,
                Vec::new(),
            )
            .await?;
        Ok(reply.content)
    }

    async fn list_models(&self) -> Result<Vec<ModelCapabilities>> {
        let url = format!("{}/models", self.base_url);
        let resp = self
            .client
            .get(&url)
            .send()
            .await
            .with_context(|| format!("GET {} failed", url))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            bail!("llama-swap list_models failed: {} — {}", status, body);
        }

        let parsed: Value = resp.json().await.context("parsing models response")?;
        let data = parsed
            .get("data")
            .and_then(|v| v.as_array())
            .ok_or_else(|| anyhow!("models response missing data[]"))?;

        let caps: Vec<ModelCapabilities> = data
            .iter()
            .map(|m| self.parse_model_capabilities(m))
            .collect();

        Ok(caps)
    }

    async fn model_capabilities(&self, model: &str) -> Result<ModelCapabilities> {
        let all = self.list_models().await?;
        all.into_iter()
            .find(|m| m.name == model)
            .ok_or_else(|| anyhow!("model '{}' not found on llama-swap", model))
    }

    fn primary_model(&self) -> &str {
        &self.primary_model
    }

    /// llama-server fixes its window at launch (`-c`), so ask it via
    /// llama-swap's `/upstream/<model>/props`. Cached per model; see
    /// [`LlamaCppClient::served_context_window`].
    async fn context_window(&self) -> Option<usize> {
        self.served_context_window(&self.primary_model).await
    }
}

impl LlamaCppClient {
    fn parse_model_capabilities(&self, m: &Value) -> ModelCapabilities {
        let name = m
            .get("id")
            .and_then(|v| v.as_str())
            .unwrap_or_default()
            .to_string();
        // llama-swap doesn't expose per-model modality flags. Assume all
        // configured models support vision and tool calling; a model
        // without multimodal support surfaces as a chat-call error.
        ModelCapabilities {
            name,
            has_vision: true,
            has_tool_calling: true,
            context_window: None,
        }
    }
}

/// Extract a diagnostic fragment from a llama-swap / llama-server response
/// that doesn't match the expected `{choices: [...]}` shape. llama-server
/// returns errors as `{"error": {"message": "...", "code": N, "type": "..."}}`;
fn log_timings(parsed: &Value, prompt_tokens: Option<i32>, completion_tokens: Option<i32>) {
    let timings = match parsed.get("timings") {
        Some(t) => t,
        None => return,
    };
    let prompt_tps = timings.get("prompt_per_second").and_then(|v| v.as_f64());
    let gen_tps = timings.get("predicted_per_second").and_then(|v| v.as_f64());
    let prompt_ms = timings.get("prompt_ms").and_then(|v| v.as_f64());
    let gen_ms = timings.get("predicted_ms").and_then(|v| v.as_f64());

    let mut parts: Vec<String> = Vec::new();
    if let Some(c) = prompt_tokens {
        let mut s = format!("prompt={} tok", c);
        if let Some(ms) = prompt_ms {
            s.push_str(&format!(" ({:.0} ms", ms));
            if let Some(tps) = prompt_tps {
                s.push_str(&format!(", {:.1} tok/s", tps));
            }
            s.push(')');
        }
        parts.push(s);
    }
    if let Some(c) = completion_tokens {
        let mut s = format!("gen={} tok", c);
        if let Some(ms) = gen_ms {
            s.push_str(&format!(" ({:.0} ms", ms));
            if let Some(tps) = gen_tps {
                s.push_str(&format!(", {:.1} tok/s", tps));
            }
            s.push(')');
        }
        parts.push(s);
    }
    if !parts.is_empty() {
        log::info!("llama-swap chat metrics — {}", parts.join(", "));
    }
}

/// llama-swap itself sometimes wraps subprocess failures with its own
/// `{"error": "..."}` flat shape. Surface either when present, otherwise fall
/// back to a truncated raw-JSON view.
fn extract_error_detail(parsed: &Value) -> String {
    if let Some(err) = parsed.get("error") {
        match err {
            Value::Object(_) => {
                let message = err
                    .get("message")
                    .and_then(|v| v.as_str())
                    .unwrap_or("(no message)");
                let code = err
                    .get("code")
                    .map(|v| match v {
                        Value::String(s) => s.clone(),
                        other => other.to_string(),
                    })
                    .unwrap_or_else(|| "?".to_string());
                let short_message: String = message.chars().take(240).collect();
                return format!("error code={} message=\"{}\"", code, short_message);
            }
            Value::String(s) => {
                let short: String = s.chars().take(240).collect();
                return format!("error=\"{}\"", short);
            }
            _ => {}
        }
    }
    let raw = parsed.to_string();
    raw.chars().take(300).collect()
}

/// Per-index assembly state for streamed OpenAI-style tool-call deltas:
/// `index → (id, name, concatenated argument fragments)`. BTreeMap so the
/// finalized calls come out in index order.
type ToolCallAssembly = std::collections::BTreeMap<usize, (Option<String>, Option<String>, String)>;

/// Fold one SSE frame's `delta.tool_calls` array into the assembly state.
/// Deltas carrying the same `index` merge into one call (llama.cpp streams a
/// call's argument JSON in fragments — they concatenate); distinct indexes
/// accumulate as separate calls.
fn apply_tool_call_deltas(state: &mut ToolCallAssembly, tcs: &[Value]) {
    for tc_delta in tcs {
        let idx = tc_delta.get("index").and_then(|n| n.as_u64()).unwrap_or(0) as usize;
        let entry = state.entry(idx).or_insert((None, None, String::new()));
        if let Some(id) = tc_delta.get("id").and_then(|v| v.as_str()) {
            entry.0 = Some(id.to_string());
        }
        if let Some(func) = tc_delta.get("function") {
            if let Some(name) = func.get("name").and_then(|v| v.as_str()) {
                entry.1 = Some(name.to_string());
            }
            if let Some(args) = func.get("arguments").and_then(|v| v.as_str()) {
                entry.2.push_str(args);
            }
        }
    }
}

/// Convert assembled tool-call state into canonical `ToolCall`s, parsing each
/// call's concatenated argument JSON (empty / malformed → `{}`). `None` when
/// no tool-call deltas arrived.
fn finalize_tool_calls(state: ToolCallAssembly) -> Option<Vec<ToolCall>> {
    if state.is_empty() {
        return None;
    }
    let mut v = Vec::with_capacity(state.len());
    for (_idx, (id, name, args)) in state {
        let arguments: Value = if args.trim().is_empty() {
            Value::Object(Default::default())
        } else {
            serde_json::from_str(&args).unwrap_or_else(|_| Value::Object(Default::default()))
        };
        v.push(ToolCall {
            id,
            function: ToolCallFunction {
                name: name.unwrap_or_default(),
                arguments,
            },
        });
    }
    Some(v)
}

fn find_double_newline(buf: &[u8]) -> Option<usize> {
    for i in 0..buf.len().saturating_sub(1) {
        if buf[i] == b'\n' && buf[i + 1] == b'\n' {
            return Some(i);
        }
        if i + 3 < buf.len()
            && buf[i] == b'\r'
            && buf[i + 1] == b'\n'
            && buf[i + 2] == b'\r'
            && buf[i + 3] == b'\n'
        {
            return Some(i + 1);
        }
    }
    None
}

fn image_to_data_url(img: &str) -> String {
    if img.starts_with("data:") {
        img.to_string()
    } else {
        format!("data:image/jpeg;base64,{}", img)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// One-shot HTTP server answering any request with `body`; returns its
    /// base URL and a handle yielding the request line it received.
    async fn serve_once(body: &'static str) -> (String, tokio::task::JoinHandle<String>) {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let handle = tokio::spawn(async move {
            let (mut sock, _) = listener.accept().await.unwrap();
            let mut buf = vec![0u8; 4096];
            let n = sock.read(&mut buf).await.unwrap();
            let req = String::from_utf8_lossy(&buf[..n])
                .lines()
                .next()
                .unwrap_or("")
                .to_string();
            let resp = format!(
                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
                body.len(),
                body
            );
            sock.write_all(resp.as_bytes()).await.unwrap();
            req
        });
        (format!("http://{addr}/v1"), handle)
    }

    #[tokio::test]
    async fn served_context_window_reads_props_through_upstream() {
        let (base, req) =
            serve_once(r#"{"default_generation_settings":{"n_ctx":32768},"total_slots":1}"#).await;
        let client = LlamaCppClient::new(Some(base), Some("gemma4-31b-qat-uncensored".into()));
        assert_eq!(client.context_window().await, Some(32768));
        assert_eq!(
            req.await.unwrap(),
            "GET /upstream/gemma4-31b-qat-uncensored/props HTTP/1.1"
        );
        // Cached: a second lookup makes no request (the server is gone).
        assert_eq!(client.context_window().await, Some(32768));
    }

    #[tokio::test]
    async fn served_context_window_is_none_when_unreachable() {
        let client = LlamaCppClient::new(Some("http://127.0.0.1:9/v1".into()), Some("nope".into()));
        assert_eq!(client.context_window().await, None);
    }

    #[test]
    fn tool_call_arguments_stringified_on_send() {
        let msg = ChatMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: Some(vec![ToolCall {
                id: Some("call_abc".into()),
                function: ToolCallFunction {
                    name: "search_sms".into(),
                    arguments: json!({"query": "hello", "limit": 5}),
                },
            }]),
            images: None,
        };

        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        let tcs = wire[0]
            .get("tool_calls")
            .and_then(|v| v.as_array())
            .expect("tool_calls present");
        let args = tcs[0]
            .get("function")
            .and_then(|f| f.get("arguments"))
            .and_then(|a| a.as_str())
            .expect("arguments stringified");
        let parsed: Value = serde_json::from_str(args).unwrap();
        assert_eq!(parsed["query"], "hello");
        assert_eq!(parsed["limit"], 5);
    }

    #[test]
    fn tool_call_arguments_parsed_on_receive() {
        let response_msg = json!({
            "role": "assistant",
            "content": "",
            "tool_calls": [{
                "id": "call_xyz",
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "arguments": "{\"city\":\"Boston\",\"units\":\"celsius\"}"
                }
            }]
        });
        let parsed = LlamaCppClient::openai_message_to_chat(&response_msg).unwrap();
        let tcs = parsed.tool_calls.unwrap();
        assert_eq!(tcs.len(), 1);
        assert_eq!(tcs[0].function.name, "get_weather");
        assert_eq!(tcs[0].function.arguments["city"], "Boston");
        assert_eq!(tcs[0].function.arguments["units"], "celsius");
        assert_eq!(tcs[0].id.as_deref(), Some("call_xyz"));
    }

    #[test]
    fn tool_call_arguments_accept_native_json_on_receive() {
        // Some llama.cpp builds emit arguments as a JSON object directly when
        // jinja's tool-output strict-string rule isn't applied — accept both.
        let response_msg = json!({
            "role": "assistant",
            "content": "",
            "tool_calls": [{
                "id": "call_1",
                "type": "function",
                "function": {
                    "name": "foo",
                    "arguments": {"nested": {"k": 1}}
                }
            }]
        });
        let parsed = LlamaCppClient::openai_message_to_chat(&response_msg).unwrap();
        let tc = &parsed.tool_calls.unwrap()[0];
        assert_eq!(tc.function.arguments["nested"]["k"], 1);
    }

    #[test]
    fn images_become_content_parts() {
        let mut msg = ChatMessage::user("What is in this photo?");
        msg.images = Some(vec!["BASE64DATA".into()]);

        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        let content = wire[0].get("content").and_then(|v| v.as_array()).unwrap();
        assert_eq!(content.len(), 2);
        assert_eq!(content[0]["type"], "text");
        assert_eq!(content[0]["text"], "What is in this photo?");
        assert_eq!(content[1]["type"], "image_url");
        assert_eq!(
            content[1]["image_url"]["url"],
            "data:image/jpeg;base64,BASE64DATA"
        );
    }

    #[test]
    fn data_url_images_pass_through_unchanged() {
        let mut msg = ChatMessage::user("");
        msg.images = Some(vec!["data:image/png;base64,ABCDEF".into()]);
        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        let content = wire[0].get("content").and_then(|v| v.as_array()).unwrap();
        assert_eq!(content.len(), 1);
        assert_eq!(
            content[0]["image_url"]["url"],
            "data:image/png;base64,ABCDEF"
        );
    }

    #[test]
    fn text_only_message_stays_string() {
        let msg = ChatMessage::user("hello");
        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        assert_eq!(wire[0]["content"], "hello");
        assert!(wire[0]["content"].as_str().is_some());
    }

    #[test]
    fn tool_result_inherits_tool_call_id_from_prior_assistant() {
        let assistant = ChatMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: Some(vec![ToolCall {
                id: Some("call_42".into()),
                function: ToolCallFunction {
                    name: "lookup".into(),
                    arguments: json!({}),
                },
            }]),
            images: None,
        };
        let tool_result = ChatMessage::tool_result("found it");

        let wire = LlamaCppClient::messages_to_openai(&[assistant, tool_result]);
        assert_eq!(wire[1]["role"], "tool");
        assert_eq!(wire[1]["tool_call_id"], "call_42");
    }

    #[test]
    fn multiple_tool_results_map_to_sequential_call_ids() {
        let assistant = ChatMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: Some(vec![
                ToolCall {
                    id: Some("call_A".into()),
                    function: ToolCallFunction {
                        name: "a".into(),
                        arguments: json!({}),
                    },
                },
                ToolCall {
                    id: Some("call_B".into()),
                    function: ToolCallFunction {
                        name: "b".into(),
                        arguments: json!({}),
                    },
                },
            ]),
            images: None,
        };
        let r1 = ChatMessage::tool_result("a result");
        let r2 = ChatMessage::tool_result("b result");

        let wire = LlamaCppClient::messages_to_openai(&[assistant, r1, r2]);
        assert_eq!(wire[1]["tool_call_id"], "call_A");
        assert_eq!(wire[2]["tool_call_id"], "call_B");
    }

    #[test]
    fn missing_tool_call_id_gets_synthetic_fallback() {
        let assistant = ChatMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: Some(vec![ToolCall {
                id: None,
                function: ToolCallFunction {
                    name: "noid".into(),
                    arguments: json!({}),
                },
            }]),
            images: None,
        };
        let wire = LlamaCppClient::messages_to_openai(&[assistant]);
        let tcs = wire[0]
            .get("tool_calls")
            .and_then(|v| v.as_array())
            .unwrap();
        assert_eq!(tcs[0]["id"], "call_0");
    }

    #[test]
    fn all_models_report_vision_capable() {
        let mut c = LlamaCppClient::new(None, Some("chat".into()));
        c.set_vision_model("vision".into());

        let m_chat = json!({ "id": "chat" });
        let m_vision = json!({ "id": "vision" });
        let m_other = json!({ "id": "embed" });

        let chat = c.parse_model_capabilities(&m_chat);
        let vision = c.parse_model_capabilities(&m_vision);
        let other = c.parse_model_capabilities(&m_other);

        assert!(chat.has_vision);
        assert!(chat.has_tool_calling);
        assert!(vision.has_vision);
        assert!(other.has_vision);
    }

    #[test]
    fn vision_model_defaults_to_primary() {
        let c = LlamaCppClient::new(None, Some("qwen3:32b".into()));
        assert_eq!(c.primary_model, "qwen3:32b");
        assert_eq!(c.vision_model, "qwen3:32b");
    }

    #[test]
    fn vision_model_explicit_override_diverges_from_primary() {
        let mut c = LlamaCppClient::new(None, Some("qwen3:32b".into()));
        c.set_vision_model("minicpm-v".into());
        assert_eq!(c.primary_model, "qwen3:32b");
        assert_eq!(c.vision_model, "minicpm-v");
    }

    #[test]
    fn cloned_local_with_model_override_keeps_all_slots_consistent() {
        // Simulates what resolve_backend does for the `local` client:
        // clone the configured client, then override primary + vision
        // to match the user-selected chat model. This prevents mid-turn
        // model swaps in llama-swap exclusive mode.
        let mut base = LlamaCppClient::new(None, Some("chat".into()));
        base.set_vision_model("vision".into());
        base.set_embedding_model("embed".into());

        let mut local = base.clone();
        let user_selected = "qwen3:32b";
        local.primary_model = user_selected.to_string();
        local.set_vision_model(user_selected.to_string());

        // Chat generation (rerank) routes through primary_model.
        assert_eq!(local.primary_model, user_selected);
        // describe_image routes through vision_model.
        assert_eq!(local.vision_model, user_selected);
        // Embeddings stay on the dedicated slot — separate endpoint,
        // no model swap conflict.
        assert_eq!(local.embedding_model, "embed");
    }

    #[test]
    fn hybrid_mode_local_client_preserves_vision_model() {
        // In hybrid mode, overrides.model is an OpenRouter model id
        // (e.g. "google/gemini-3-flash-preview"). The local llamacpp
        // client must NOT adopt that as its vision_model — it should
        // keep the configured LLAMA_SWAP_VISION_MODEL so describe_image
        // hits the correct local slot instead of sending an unknown
        // model name to llama-swap.
        let mut base = LlamaCppClient::new(None, Some("chat".into()));
        base.set_vision_model("vision".into());
        base.set_embedding_model("embed".into());

        // Simulate what resolve_backend SHOULD do for hybrid mode:
        // clone but do NOT override primary_model / vision_model.
        let local = base.clone();

        // The local client keeps its configured slots.
        assert_eq!(local.primary_model, "chat");
        assert_eq!(local.vision_model, "vision");
        assert_eq!(local.embedding_model, "embed");
    }

    #[test]
    fn assistant_tool_calls_emit_null_content() {
        let msg = ChatMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: Some(vec![ToolCall {
                id: Some("call_1".into()),
                function: ToolCallFunction {
                    name: "search".into(),
                    arguments: json!({}),
                },
            }]),
            images: None,
        };
        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        assert!(
            wire[0]["content"].is_null(),
            "empty content + tool_calls should emit null"
        );
    }

    #[test]
    fn assistant_with_content_and_tool_calls_preserves_content() {
        let msg = ChatMessage {
            role: "assistant".into(),
            content: "Let me search for that.".into(),
            tool_calls: Some(vec![ToolCall {
                id: Some("call_1".into()),
                function: ToolCallFunction {
                    name: "search".into(),
                    arguments: json!({}),
                },
            }]),
            images: None,
        };
        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        assert_eq!(wire[0]["content"], "Let me search for that.");
    }

    #[test]
    fn assistant_without_tool_calls_keeps_empty_string_content() {
        let msg = ChatMessage {
            role: "assistant".into(),
            content: String::new(),
            tool_calls: None,
            images: None,
        };
        let wire = LlamaCppClient::messages_to_openai(&[msg]);
        assert_eq!(wire[0]["content"], "");
    }

    #[test]
    fn swap_root_strips_v1_suffix() {
        let c = LlamaCppClient::new(Some("http://localhost:9292/v1".to_string()), None);
        assert_eq!(c.swap_root(), "http://localhost:9292");

        // Tolerates a trailing slash on the base URL.
        let c2 = LlamaCppClient::new(Some("http://localhost:9292/v1/".to_string()), None);
        assert_eq!(c2.swap_root(), "http://localhost:9292");

        // No /v1 suffix → returned unchanged.
        let c3 = LlamaCppClient::new(Some("http://host:1234".to_string()), None);
        assert_eq!(c3.swap_root(), "http://host:1234");
    }

    #[test]
    fn tts_model_defaults_to_chatterbox() {
        let c = LlamaCppClient::new(None, None);
        assert_eq!(c.tts_model, "chatterbox");
    }

    #[test]
    fn stream_assembly_keeps_two_tool_calls_from_separate_chunks() {
        // llama.cpp emits one delta per SSE frame; two calls with distinct
        // `index` values arriving in separate frames must BOTH survive.
        let mut state = ToolCallAssembly::new();
        apply_tool_call_deltas(
            &mut state,
            &[json!({
                "index": 0,
                "id": "call_a",
                "function": { "name": "get_sms_messages", "arguments": "{\"date\":\"2019-01-01\"}" }
            })],
        );
        apply_tool_call_deltas(
            &mut state,
            &[json!({
                "index": 1,
                "id": "call_b",
                "function": { "name": "reverse_geocode", "arguments": "{\"latitude\":1.0,\"longitude\":2.0}" }
            })],
        );

        let calls = finalize_tool_calls(state).expect("two calls assembled");
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].id.as_deref(), Some("call_a"));
        assert_eq!(calls[0].function.name, "get_sms_messages");
        assert_eq!(calls[0].function.arguments["date"], "2019-01-01");
        assert_eq!(calls[1].id.as_deref(), Some("call_b"));
        assert_eq!(calls[1].function.name, "reverse_geocode");
        assert_eq!(calls[1].function.arguments["latitude"], 1.0);
    }

    #[test]
    fn stream_assembly_concatenates_argument_fragments_for_same_index() {
        // A single call's argument JSON streamed across frames concatenates
        // into one parseable document.
        let mut state = ToolCallAssembly::new();
        apply_tool_call_deltas(
            &mut state,
            &[json!({
                "index": 0,
                "id": "call_x",
                "function": { "name": "search_messages", "arguments": "{\"query\":" }
            })],
        );
        apply_tool_call_deltas(
            &mut state,
            &[json!({
                "index": 0,
                "function": { "arguments": "\"dinner\"}" }
            })],
        );

        let calls = finalize_tool_calls(state).expect("one call assembled");
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].function.name, "search_messages");
        assert_eq!(calls[0].function.arguments["query"], "dinner");
    }

    #[test]
    fn stream_assembly_empty_state_finalizes_to_none() {
        assert!(finalize_tool_calls(ToolCallAssembly::new()).is_none());
    }
}