octolib 0.4.2

// Copyright 2025 Muvon Un Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Jina AI embedding provider implementation

use anyhow::{Context, Result};
use serde_json::{json, Value};

use super::super::types::InputType;
use super::{EmbeddingProvider, HTTP_CLIENT};

/// Jina provider implementation for trait
pub struct JinaProviderImpl {
    model_name: String,
    dimension: usize,
}

impl JinaProviderImpl {
    pub fn new(model: &str) -> Result<Self> {
        // Validate model first - fail fast if unsupported
        let supported_models = [
            "jina-embeddings-v4",
            "jina-clip-v2",
            "jina-embeddings-v3",
            "jina-clip-v1",
            "jina-embeddings-v2-base-es",
            "jina-embeddings-v2-base-code",
            "jina-embeddings-v2-base-de",
            "jina-embeddings-v2-base-zh",
            "jina-embeddings-v2-base-en",
        ];

        if !supported_models.contains(&model) {
            return Err(anyhow::anyhow!(
                "Unsupported Jina model: '{}'. Supported models: {:?}",
                model,
                supported_models
            ));
        }

        let dimension = Self::get_model_dimension(model);
        Ok(Self {
            model_name: model.to_string(),
            dimension,
        })
    }

    fn get_model_dimension(model: &str) -> usize {
        match model {
            "jina-embeddings-v4" => 2048,
            "jina-clip-v2" => 1024,
            "jina-embeddings-v3" => 1024,
            "jina-clip-v1" => 768,
            "jina-embeddings-v2-base-es" => 768,
            "jina-embeddings-v2-base-code" => 768,
            "jina-embeddings-v2-base-de" => 768,
            "jina-embeddings-v2-base-zh" => 768,
            "jina-embeddings-v2-base-en" => 768,
            _ => {
                // This should never be reached due to validation in new()
                panic!(
                    "Invalid Jina model '{}' passed to get_model_dimension",
                    model
                );
            }
        }
    }
}

#[async_trait::async_trait]
impl EmbeddingProvider for JinaProviderImpl {
    async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>> {
        JinaProvider::generate_embeddings(text, &self.model_name).await
    }

    async fn generate_embeddings_batch(
        &self,
        texts: Vec<String>,
        input_type: InputType,
    ) -> Result<Vec<Vec<f32>>> {
        // Apply prefix manually for Jina (doesn't support input_type API)
        let processed_texts: Vec<String> = texts
            .into_iter()
            .map(|text| input_type.apply_prefix(&text))
            .collect();
        JinaProvider::generate_embeddings_batch(processed_texts, &self.model_name).await
    }

    fn get_dimension(&self) -> usize {
        self.dimension
    }

    fn is_model_supported(&self) -> bool {
        // REAL validation - only support actual Jina models
        matches!(
            self.model_name.as_str(),
            "jina-embeddings-v4"
                | "jina-clip-v2"
                | "jina-embeddings-v3"
                | "jina-clip-v1"
                | "jina-embeddings-v2-base-es"
                | "jina-embeddings-v2-base-code"
                | "jina-embeddings-v2-base-de"
                | "jina-embeddings-v2-base-zh"
                | "jina-embeddings-v2-base-en"
        )
    }
}

/// Jina provider implementation
pub struct JinaProvider;

impl JinaProvider {
    pub async fn generate_embeddings(contents: &str, model: &str) -> Result<Vec<f32>> {
        let result = Self::generate_embeddings_batch(vec![contents.to_string()], model).await?;
        result
            .first()
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("No embeddings found"))
    }

    pub async fn generate_embeddings_batch(
        texts: Vec<String>,
        model: &str,
    ) -> Result<Vec<Vec<f32>>> {
        let jina_api_key =
            std::env::var("JINA_API_KEY").context("JINA_API_KEY environment variable not set")?;

        let response = HTTP_CLIENT
            .post("https://api.jina.ai/v1/embeddings")
            .header("Authorization", format!("Bearer {}", jina_api_key))
            .json(&json!({
                "input": texts,
                "model": model,
            }))
            .send()
            .await?;

        let response_json: Value = response.json().await?;

        let embeddings = response_json["data"]
            .as_array()
            .context("Failed to get embeddings array")?
            .iter()
            .map(|data| {
                data["embedding"]
                    .as_array()
                    .unwrap_or(&Vec::new())
                    .iter()
                    .map(|v| v.as_f64().unwrap_or_default() as f32)
                    .collect()
            })
            .collect();

        Ok(embeddings)
    }
}