mofa-foundation 0.1.1

//! Google Gemini Provider (experimental, text-only)
//!
//! Implements Gemini Pro via Generative Language API v1beta.

use super::provider::{ChatStream, LLMProvider, ModelCapabilities, ModelInfo};
use super::types::*;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Gemini provider configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeminiConfig {
    /// API key
    pub api_key: String,
    /// Base URL (default: https://generativelanguage.googleapis.com)
    pub base_url: String,
    /// Default model id, e.g., gemini-1.5-pro-latest
    pub default_model: String,
    /// Default temperature
    pub default_temperature: f32,
    /// Default max output tokens
    pub default_max_tokens: u32,
    /// Request timeout
    pub timeout_secs: u64,
}

impl Default for GeminiConfig {
    fn default() -> Self {
        Self {
            api_key: String::new(),
            base_url: "https://generativelanguage.googleapis.com".to_string(),
            default_model: "gemini-1.5-pro-latest".to_string(),
            default_temperature: 0.7,
            default_max_tokens: 2048,
            timeout_secs: 60,
        }
    }
}

impl GeminiConfig {
    pub fn new(api_key: impl Into<String>) -> Self {
        Self {
            api_key: api_key.into(),
            ..Default::default()
        }
    }

    pub fn from_env() -> Self {
        let mut cfg = Self {
            api_key: std::env::var("GEMINI_API_KEY").unwrap_or_default(),
            ..Default::default()
        };

        if let Ok(model) = std::env::var("GEMINI_MODEL") {
            cfg.default_model = model;
        }
        if let Ok(base_url) = std::env::var("GEMINI_BASE_URL") {
            cfg.base_url = base_url;
        }
        cfg
    }

    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = url.into();
        self
    }

    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.default_model = model.into();
        self
    }

    pub fn with_temperature(mut self, temp: f32) -> Self {
        self.default_temperature = temp;
        self
    }

    pub fn with_max_tokens(mut self, tokens: u32) -> Self {
        self.default_max_tokens = tokens;
        self
    }

    pub fn with_timeout(mut self, secs: u64) -> Self {
        self.timeout_secs = secs;
        self
    }
}

/// Gemini provider (text-only, no tools for now)
pub struct GeminiProvider {
    client: reqwest::Client,
    config: GeminiConfig,
}

impl GeminiProvider {
    pub fn new(api_key: impl Into<String>) -> Self {
        Self::with_config(GeminiConfig::new(api_key))
    }

    pub fn from_env() -> Self {
        Self::with_config(GeminiConfig::from_env())
    }

    pub fn with_config(config: GeminiConfig) -> Self {
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(config.timeout_secs))
            .build()
            .expect("Failed to build reqwest client");
        Self { client, config }
    }

    fn convert_messages(messages: &[ChatMessage]) -> (Option<String>, Vec<serde_json::Value>) {
        let mut system_parts = Vec::new();
        let mut contents = Vec::new();

        for msg in messages {
            match msg.role {
                Role::System => {
                    if let Some(text) = msg.text_content() {
                        system_parts.push(text.to_string());
                    }
                }
                Role::User | Role::Assistant | Role::Tool => {
                    let role = match msg.role {
                        Role::User => "user",
                        Role::Assistant => "model",
                        Role::Tool => "user",
                        Role::System => unreachable!(),
                    };

                    let text = match &msg.content {
                        Some(MessageContent::Text(t)) => t.clone(),
                        Some(MessageContent::Parts(parts)) => parts
                            .iter()
                            .filter_map(|p| match p {
                                ContentPart::Text { text } => Some(text.clone()),
                                _ => None,
                            })
                            .collect::<Vec<_>>()
                            .join("\n"),
                        None => String::new(),
                    };

                    contents.push(serde_json::json!({
                        "role": role,
                        "parts": [{"text": text}],
                    }));
                }
            }
        }

        let system = if system_parts.is_empty() {
            None
        } else {
            Some(system_parts.join("\n"))
        };

        (system, contents)
    }

    fn map_error(err: reqwest::Error) -> LLMError {
        if err.is_timeout() {
            LLMError::Timeout(err.to_string())
        } else if err.is_connect() || err.is_request() {
            LLMError::NetworkError(err.to_string())
        } else {
            LLMError::Other(err.to_string())
        }
    }
}

#[derive(Debug, Deserialize)]
struct GeminiCandidate {
    content: Option<GeminiContent>,
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
struct GeminiContent {
    parts: Vec<GeminiPart>,
}

#[derive(Debug, Deserialize)]
struct GeminiPart {
    text: Option<String>,
}

#[derive(Debug, Deserialize)]
struct GeminiUsage {
    #[serde(rename = "promptTokenCount")]
    prompt_tokens: u32,
    #[serde(rename = "candidatesTokenCount")]
    candidates_tokens: u32,
    #[serde(rename = "totalTokenCount")]
    total_tokens: u32,
}

#[derive(Debug, Deserialize)]
struct GeminiResponse {
    candidates: Vec<GeminiCandidate>,
    #[serde(rename = "usageMetadata")]
    usage: Option<GeminiUsage>,
    model_version: Option<String>,
}

#[async_trait]
impl LLMProvider for GeminiProvider {
    fn name(&self) -> &str {
        "gemini"
    }

    fn default_model(&self) -> &str {
        &self.config.default_model
    }

    fn supports_streaming(&self) -> bool {
        false
    }

    fn supports_tools(&self) -> bool {
        false
    }

    fn supports_vision(&self) -> bool {
        false
    }

    async fn chat(&self, request: ChatCompletionRequest) -> LLMResult<ChatCompletionResponse> {
        let (system, contents) = Self::convert_messages(&request.messages);

        let model = if request.model.is_empty() {
            self.config.default_model.clone()
        } else {
            request.model.clone()
        };

        let mut body = serde_json::json!({
            "contents": contents,
            "generationConfig": {
                "temperature": request
                    .temperature
                    .unwrap_or(self.config.default_temperature),
                "maxOutputTokens": request
                    .max_tokens
                    .unwrap_or(self.config.default_max_tokens),
            }
        });

        if let Some(sys) = system {
            body["systemInstruction"] = serde_json::json!({"parts": [{"text": sys}]});
        }

        if let Some(tp) = request.top_p {
            body["generationConfig"]["topP"] = serde_json::json!(tp);
        }

        if let Some(stop) = request.stop.clone() {
            body["generationConfig"]["stopSequences"] = serde_json::json!(stop);
        }

        let url = format!(
            "{}/v1beta/models/{}:generateContent?key={}",
            self.config.base_url.trim_end_matches('/'),
            model,
            self.config.api_key
        );

        let resp = self
            .client
            .post(&url)
            .json(&body)
            .send()
            .await
            .map_err(Self::map_error)?;

        let status = resp.status();
        let text = resp.text().await.map_err(Self::map_error)?;

        if !status.is_success() {
            return Err(LLMError::ApiError {
                code: Some(status.as_u16().to_string()),
                message: text,
            });
        }

        let parsed: GeminiResponse =
            serde_json::from_str(&text).map_err(|e| LLMError::SerializationError(e.to_string()))?;

        let first = parsed
            .candidates
            .into_iter()
            .find_map(|c| c.content)
            .and_then(|c| c.parts.into_iter().find_map(|p| p.text))
            .unwrap_or_default();

        let finish_reason = None; // Gemini returns fine-grained reasons; map later if needed

        let usage = parsed.usage.map(|u| Usage {
            prompt_tokens: u.prompt_tokens,
            completion_tokens: u.candidates_tokens,
            total_tokens: u.total_tokens,
        });

        let choice = Choice {
            index: 0,
            message: ChatMessage {
                role: Role::Assistant,
                content: Some(MessageContent::Text(first)),
                name: None,
                tool_calls: None,
                tool_call_id: None,
            },
            finish_reason,
            logprobs: None,
        };

        let created = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0))
            .as_secs();

        Ok(ChatCompletionResponse {
            id: "".to_string(),
            object: "chat.completion".to_string(),
            created,
            model,
            choices: vec![choice],
            usage,
            system_fingerprint: parsed.model_version,
        })
    }

    async fn chat_stream(&self, _request: ChatCompletionRequest) -> LLMResult<ChatStream> {
        Err(LLMError::ProviderNotSupported(
            "Gemini streaming not yet implemented".to_string(),
        ))
    }

    async fn health_check(&self) -> LLMResult<bool> {
        let req = ChatCompletionRequest::new(self.default_model())
            .system("Say 'ok'")
            .max_tokens(4);
        self.chat(req).await.map(|_| true).or(Ok(false))
    }

    async fn get_model_info(&self, model: &str) -> LLMResult<ModelInfo> {
        Ok(ModelInfo {
            id: model.to_string(),
            name: model.to_string(),
            description: None,
            context_window: None,
            max_output_tokens: Some(self.config.default_max_tokens),
            training_cutoff: None,
            capabilities: ModelCapabilities {
                streaming: false,
                tools: false,
                vision: false,
                json_mode: true,
                json_schema: false,
            },
        })
    }
}