coderlib 0.1.0

//! Google Gemini provider implementation
//!
//! This module provides native integration with Google's Gemini AI models
//! including Gemini Pro, Gemini Pro Vision, and other variants.

use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use tokio_stream::{Stream, StreamExt};
use tracing::{debug, info, error};

use crate::core::error::ProviderError;
use crate::core::TokenUsage;
use crate::llm::{Provider, ProviderResponse, ProviderEvent, Model};
use crate::storage::{Message, MessageRole, MessageContent};
use crate::tools::Tool;

/// Google Gemini provider
pub struct GeminiProvider {
    client: Client,
    api_key: String,
    base_url: String,
    model: Model,
}

/// Gemini API request structure
#[derive(Debug, Serialize)]
struct GeminiRequest {
    contents: Vec<GeminiContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    generation_config: Option<GenerationConfig>,
    #[serde(skip_serializing_if = "Option::is_none")]
    safety_settings: Option<Vec<SafetySetting>>,
}

/// Gemini content structure
#[derive(Debug, Serialize)]
struct GeminiContent {
    role: String,
    parts: Vec<GeminiPart>,
}

/// Gemini content part
#[derive(Debug, Serialize)]
struct GeminiPart {
    text: String,
}

/// Generation configuration
#[derive(Debug, Serialize)]
struct GenerationConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_k: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_output_tokens: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    candidate_count: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stop_sequences: Option<Vec<String>>,
}

/// Safety setting for content filtering
#[derive(Debug, Serialize)]
struct SafetySetting {
    category: String,
    threshold: String,
}

/// Gemini API response
#[derive(Debug, Serialize, Deserialize)]
struct GeminiResponse {
    candidates: Vec<GeminiCandidate>,
    #[serde(rename = "usageMetadata")]
    usage_metadata: Option<GeminiUsageMetadata>,
    #[serde(rename = "promptFeedback")]
    prompt_feedback: Option<GeminiPromptFeedback>,
}

/// Gemini response candidate
#[derive(Debug, Serialize, Deserialize)]
struct GeminiCandidate {
    content: GeminiResponseContent,
    #[serde(rename = "finishReason")]
    finish_reason: Option<String>,
    index: Option<i32>,
    #[serde(rename = "safetyRatings")]
    safety_ratings: Option<Vec<GeminiSafetyRating>>,
}

/// Gemini response content
#[derive(Debug, Serialize, Deserialize)]
struct GeminiResponseContent {
    parts: Vec<GeminiResponsePart>,
    role: Option<String>,
}

/// Gemini response part
#[derive(Debug, Serialize, Deserialize)]
struct GeminiResponsePart {
    text: String,
}

/// Usage metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
struct GeminiUsageMetadata {
    #[serde(rename = "promptTokenCount")]
    prompt_token_count: Option<i32>,
    #[serde(rename = "candidatesTokenCount")]
    candidates_token_count: Option<i32>,
    #[serde(rename = "totalTokenCount")]
    total_token_count: Option<i32>,
}

/// Prompt feedback
#[derive(Debug, Serialize, Deserialize)]
struct GeminiPromptFeedback {
    #[serde(rename = "blockReason")]
    block_reason: Option<String>,
    #[serde(rename = "safetyRatings")]
    safety_ratings: Option<Vec<GeminiSafetyRating>>,
}

/// Safety rating
#[derive(Debug, Serialize, Deserialize)]
struct GeminiSafetyRating {
    category: String,
    probability: String,
}

/// Streaming response chunk
#[derive(Debug, Serialize, Deserialize)]
struct GeminiStreamChunk {
    candidates: Option<Vec<GeminiCandidate>>,
    #[serde(rename = "usageMetadata")]
    usage_metadata: Option<GeminiUsageMetadata>,
}

impl GeminiProvider {
    /// Create a new Gemini provider
    pub fn new(api_key: String, model: Model) -> Result<Self, ProviderError> {
        let client = Client::new();
        let base_url = "https://generativelanguage.googleapis.com/v1beta".to_string();

        Ok(Self {
            client,
            api_key,
            base_url,
            model,
        })
    }

    /// Convert our messages to Gemini format
    fn convert_messages(&self, messages: &[Message]) -> Vec<GeminiContent> {
        messages
            .iter()
            .map(|msg| {
                let role = match msg.role {
                    MessageRole::User => "user",
                    MessageRole::Assistant => "model",
                    MessageRole::System => "user", // Gemini doesn't have system role, treat as user
                    MessageRole::Tool => "user", // Treat tool messages as user messages
                };

                let text = match &msg.content {
                    MessageContent::Text(text) => text.clone(),
                    MessageContent::Structured(parts) => {
                        // Extract text from structured content
                        parts.iter()
                            .filter_map(|part| match part {
                                crate::storage::ContentPart::Text { text } => Some(text.clone()),
                                _ => None,
                            })
                            .collect::<Vec<_>>()
                            .join(" ")
                    }
                    MessageContent::ToolCall { name, parameters, .. } => {
                        format!("Tool call: {} with parameters: {}", name, parameters)
                    }
                    MessageContent::ToolResult { result, .. } => {
                        format!("Tool result: {}", result)
                    }
                };

                GeminiContent {
                    role: role.to_string(),
                    parts: vec![GeminiPart { text }],
                }
            })
            .collect()
    }

    /// Build generation config from request parameters
    fn build_generation_config(&self, temperature: Option<f32>, max_tokens: Option<u32>) -> Option<GenerationConfig> {
        let mut config = GenerationConfig {
            temperature: None,
            top_p: None,
            top_k: None,
            max_output_tokens: None,
            candidate_count: Some(1),
            stop_sequences: None,
        };

        let mut has_config = false;

        if let Some(temp) = temperature {
            config.temperature = Some(temp);
            has_config = true;
        }

        if let Some(max_tokens) = max_tokens {
            config.max_output_tokens = Some(max_tokens as i32);
            has_config = true;
        }

        if has_config {
            Some(config)
        } else {
            None
        }
    }

    /// Build safety settings for content filtering
    fn build_safety_settings(&self) -> Vec<SafetySetting> {
        vec![
            SafetySetting {
                category: "HARM_CATEGORY_HARASSMENT".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
            SafetySetting {
                category: "HARM_CATEGORY_HATE_SPEECH".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
            SafetySetting {
                category: "HARM_CATEGORY_SEXUALLY_EXPLICIT".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
            SafetySetting {
                category: "HARM_CATEGORY_DANGEROUS_CONTENT".to_string(),
                threshold: "BLOCK_MEDIUM_AND_ABOVE".to_string(),
            },
        ]
    }

    /// Make API request to Gemini
    async fn make_request(&self, gemini_request: &GeminiRequest) -> Result<GeminiResponse, ProviderError> {
        let url = format!(
            "{}/models/{}:generateContent?key={}",
            self.base_url, self.model.id, self.api_key
        );

        debug!("Making Gemini API request to: {}", url);

        let response = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(gemini_request)
            .send()
            .await
            .map_err(|e| ProviderError::NetworkError(e.to_string()))?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            error!("Gemini API error {}: {}", status, error_text);
            return Err(ProviderError::ApiError(format!("HTTP {}: {}", status, error_text)));
        }

        let gemini_response: GeminiResponse = response
            .json()
            .await
            .map_err(|e| ProviderError::ApiError(format!("Parse error: {}", e)))?;

        Ok(gemini_response)
    }

    /// Make streaming API request to Gemini
    async fn make_streaming_request(
        &self,
        gemini_request: &GeminiRequest,
    ) -> Result<impl Stream<Item = Result<GeminiStreamChunk, ProviderError>>, ProviderError> {
        let url = format!(
            "{}/models/{}:streamGenerateContent?key={}",
            self.base_url, self.model.id, self.api_key
        );

        debug!("Making Gemini streaming API request to: {}", url);

        let response = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(gemini_request)
            .send()
            .await
            .map_err(|e| ProviderError::NetworkError(e.to_string()))?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string());
            error!("Gemini streaming API error {}: {}", status, error_text);
            return Err(ProviderError::ApiError(format!("HTTP {}: {}", status, error_text)));
        }

        let stream = response
            .bytes_stream()
            .map(|result| {
                result
                    .map_err(|e| ProviderError::NetworkError(e.to_string()))
                    .and_then(|bytes| {
                        let text = String::from_utf8_lossy(&bytes);

                        // Gemini streaming format is JSON lines
                        for line in text.lines() {
                            let line = line.trim();
                            if line.is_empty() || !line.starts_with('{') {
                                continue;
                            }

                            match serde_json::from_str::<GeminiStreamChunk>(line) {
                                Ok(chunk) => return Ok(chunk),
                                Err(e) => {
                                    debug!("Failed to parse Gemini stream chunk: {} - Line: {}", e, line);
                                    continue;
                                }
                            }
                        }

                        Err(ProviderError::ApiError("No valid JSON found in chunk".to_string()))
                    })
            })
            .filter_map(|result| {
                match result {
                    Ok(chunk) => Some(Ok(chunk)),
                    Err(e) => {
                        debug!("Skipping invalid chunk: {}", e);
                        None
                    }
                }
            });

        Ok(stream)
    }

    /// Convert Gemini response to our format
    fn convert_response(&self, gemini_response: GeminiResponse) -> Result<ProviderResponse, ProviderError> {
        if gemini_response.candidates.is_empty() {
            return Err(ProviderError::ApiError("No candidates in response".to_string()));
        }

        let candidate = &gemini_response.candidates[0];

        if candidate.content.parts.is_empty() {
            return Err(ProviderError::ApiError("No content parts in response".to_string()));
        }

        let content = candidate.content.parts
            .iter()
            .map(|part| part.text.clone())
            .collect::<Vec<_>>()
            .join("");

        let token_usage = gemini_response.usage_metadata.clone().map(|usage| TokenUsage {
            input_tokens: usage.prompt_token_count.unwrap_or(0) as u32,
            output_tokens: usage.candidates_token_count.unwrap_or(0) as u32,
            total_tokens: (usage.prompt_token_count.unwrap_or(0) + usage.candidates_token_count.unwrap_or(0)) as u32,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        });

        Ok(ProviderResponse {
            content,
            tool_calls: Vec::new(), // TODO: Implement tool calls for Gemini
            token_usage,
            metadata: serde_json::to_value(&gemini_response).unwrap_or_default(),
        })
    }
}

#[async_trait]
impl Provider for GeminiProvider {
    async fn send_messages(
        &self,
        messages: Vec<Message>,
        _tools: Vec<Box<dyn Tool>>,
    ) -> Result<ProviderResponse, ProviderError> {
        info!("Generating response with Gemini model: {}", self.model.id);

        // Convert messages to Gemini format
        let contents = self.convert_messages(&messages);

        // Build request
        let gemini_request = GeminiRequest {
            contents,
            generation_config: self.build_generation_config(Some(0.7), Some(2048)),
            safety_settings: Some(self.build_safety_settings()),
        };

        // Make API request
        let gemini_response = self.make_request(&gemini_request).await?;

        // Convert response
        self.convert_response(gemini_response)
    }

    async fn stream_response(
        &self,
        messages: Vec<Message>,
        _tools: Vec<Box<dyn Tool>>,
    ) -> Result<mpsc::Receiver<Result<ProviderEvent, ProviderError>>, ProviderError> {
        info!("Starting streaming generation with Gemini model: {}", self.model.id);

        let (tx, rx) = mpsc::channel(100);

        // Convert messages to Gemini format
        let contents = self.convert_messages(&messages);

        // Build request
        let gemini_request = GeminiRequest {
            contents,
            generation_config: self.build_generation_config(Some(0.7), Some(2048)),
            safety_settings: Some(self.build_safety_settings()),
        };

        // Clone necessary data for the async task
        let provider = self.clone();

        tokio::spawn(async move {
            match provider.make_streaming_request(&gemini_request).await {
                Ok(mut stream) => {
                    while let Some(chunk_result) = stream.next().await {
                        match chunk_result {
                            Ok(chunk) => {
                                if let Some(candidates) = chunk.candidates {
                                    if let Some(candidate) = candidates.first() {
                                        if let Some(part) = candidate.content.parts.first() {
                                            let _ = tx.send(Ok(ProviderEvent::ContentChunk {
                                                content: part.text.clone(),
                                            })).await;
                                        }

                                        if candidate.finish_reason.is_some() {
                                            let _ = tx.send(Ok(ProviderEvent::Complete {
                                                token_usage: chunk.usage_metadata.map(|usage| TokenUsage {
                                                    input_tokens: usage.prompt_token_count.unwrap_or(0) as u32,
                                                    output_tokens: usage.candidates_token_count.unwrap_or(0) as u32,
                                                    total_tokens: (usage.prompt_token_count.unwrap_or(0) + usage.candidates_token_count.unwrap_or(0)) as u32,
                                                    cache_creation_tokens: 0,
                                                    cache_read_tokens: 0,
                                                }),
                                            })).await;
                                            break;
                                        }
                                    }
                                }
                            }
                            Err(e) => {
                                let _ = tx.send(Err(e)).await;
                                break;
                            }
                        }
                    }
                }
                Err(e) => {
                    let _ = tx.send(Err(e)).await;
                }
            }
        });

        Ok(rx)
    }

    fn model(&self) -> &Model {
        &self.model
    }

    fn name(&self) -> &str {
        "gemini"
    }

    async fn is_available(&self) -> bool {
        // Simple validation - check if we have an API key
        !self.api_key.is_empty()
    }
}

impl GeminiProvider {
    /// Clone the provider
    fn clone(&self) -> Self {
        Self {
            client: self.client.clone(),
            api_key: self.api_key.clone(),
            base_url: self.base_url.clone(),
            model: self.model.clone(),
        }
    }
}