use serde::{Deserialize, Serialize};

use super::client::LlmClient;
use crate::error::AgentError;

// ── Embedding types ────────────────────────────────────────────

#[derive(Serialize)]
struct EmbeddingRequest {
    model: String,
    input: Vec<String>,
}

#[derive(Deserialize)]
struct EmbeddingResponse {
    data: Vec<EmbeddingData>,
}

#[derive(Deserialize)]
struct EmbeddingData {
    embedding: Vec<f64>,
    index: usize,
}

/// Max inputs per embedding request. The bge/OpenAI-like backends cap the input
/// array (bge-multilingual-gemma2 rejects >25 with "batch size overflow"), so we
/// chunk larger corpora — a whole control catalog (~1.8k) would otherwise 500.
const EMBED_BATCH_SIZE: usize = 16;

// ── Embedding implementation ───────────────────────────────────

impl LlmClient {
    pub fn embed_model(&self) -> &str {
        &self.embed_model
    }

    /// Generate embeddings for a batch of texts, chunking into backend-sized
    /// requests and preserving input order across chunks.
    pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }
        let mut out = Vec::with_capacity(texts.len());
        for chunk in texts.chunks(EMBED_BATCH_SIZE) {
            out.extend(self.embed_batch(chunk.to_vec()).await?);
        }
        Ok(out)
    }

    /// Embed one backend-sized batch (≤ [`EMBED_BATCH_SIZE`]) in a single request.
    async fn embed_batch(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
        let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));

        let request_body = EmbeddingRequest {
            model: self.embed_model.clone(),
            input: texts,
        };

        let mut req = self
            .http
            .post(&url)
            .header("content-type", "application/json")
            .json(&request_body);

        if let Some(auth) = self.auth_header() {
            req = req.header("Authorization", auth);
        }

        let resp = req
            .send()
            .await
            .map_err(|e| AgentError::Other(format!("Embedding request failed: {e}")))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            return Err(AgentError::Other(format!(
                "Embedding API returned {status}: {body}"
            )));
        }

        let body: EmbeddingResponse = resp
            .json()
            .await
            .map_err(|e| AgentError::Other(format!("Failed to parse embedding response: {e}")))?;

        let mut data = body.data;
        data.sort_by_key(|d| d.index);

        Ok(data.into_iter().map(|d| d.embedding).collect())
    }
}

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

    fn client() -> LlmClient {
        LlmClient::new(
            "http://unused".into(),
            SecretString::from(String::new()),
            "m".into(),
            "e".into(),
        )
    }

    #[tokio::test]
    async fn empty_input_makes_no_request() {
        // Must short-circuit before any HTTP call (base_url is unroutable).
        let out = client().embed(Vec::new()).await.unwrap();
        assert!(out.is_empty());
    }

    #[test]
    fn batch_size_is_within_backend_cap() {
        assert!(
            EMBED_BATCH_SIZE <= 25,
            "must stay under the bge 25-input cap"
        );
    }
}