use rig::client::EmbeddingsClient;
use rig::embeddings::EmbeddingModel;

use crate::{
    client::AiClient,
    error::{AiError, AiResult},
};

#[derive(Clone)]
pub struct EmbedClient {
    model_name: String,
    client: rig::providers::openai::Client,
}

impl EmbedClient {
    pub fn new(ai_client: &AiClient) -> AiResult<Self> {
        Ok(Self {
            model_name: ai_client.embed_model().to_string(),
            client: ai_client.embed_client().clone(),
        })
    }

    fn embedding_model(&self) -> impl EmbeddingModel + '_ {
        self.client.embedding_model(&self.model_name)
    }

    pub async fn embed_text(&self, text: String) -> AiResult<Vec<f32>> {
        let model = self.embedding_model();
        let mut embeddings = model
            .embed_texts(vec![text])
            .await
            .map_err(|e| AiError::Api(e.to_string()))?;
        embeddings
            .pop()
            .map(|e| e.vec.into_iter().map(|v| v as f32).collect())
            .ok_or_else(|| {
                AiError::Response("no embedding returned".to_string())
            })
    }

    pub async fn embed_texts(
        &self,
        texts: Vec<String>,
    ) -> AiResult<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }
        let model = self.embedding_model();
        let embeddings = model
            .embed_texts(texts)
            .await
            .map_err(|e| AiError::Api(e.to_string()))?;
        Ok(embeddings
            .into_iter()
            .map(|e| e.vec.into_iter().map(|v| v as f32).collect())
            .collect())
    }

    pub async fn embed_texts_chunked(
        &self,
        texts: Vec<String>,
        batch_size: usize,
    ) -> AiResult<Vec<Vec<f32>>> {
        if batch_size == 0 {
            return Err(AiError::Config("batch_size must be > 0".to_string()));
        }
        let mut embeddings: Vec<Vec<f32>> = Vec::with_capacity(texts.len());
        for chunk in texts.chunks(batch_size) {
            let model = self.embedding_model();
            let chunk_embeddings = model
                .embed_texts(chunk.to_vec())
                .await
                .map_err(|e| AiError::Api(e.to_string()))?;
            embeddings.extend(
                chunk_embeddings
                    .into_iter()
                    .map(|e| e.vec.into_iter().map(|v| v as f32).collect()),
            );
        }
        Ok(embeddings)
    }
}

pub trait AiClientEmbedExt {
    fn embedder(&self) -> AiResult<EmbedClient>;
}

impl AiClientEmbedExt for AiClient {
    fn embedder(&self) -> AiResult<EmbedClient> {
        EmbedClient::new(self)
    }
}
