//! Google Gemini API client implementation

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::env;

use super::provider::{
    LlmProvider, LlmRequest, LlmResponse, LlmStream, MessageRole, StreamChunk, StreamMetadata,
    TokenUsage,
};
use crate::error::{Error, Result};
use futures::stream::StreamExt;

const GEMINI_API_BASE: &str = "https://generativelanguage.googleapis.com/v1beta/models";

/// Gemini API client
pub struct GeminiProvider {
    api_key: String,
    client: reqwest::Client,
}

/// Gemini API request format
#[derive(Debug, Serialize)]
struct GeminiRequest {
    contents: Vec<GeminiContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(rename = "generationConfig")]
    generation_config: Option<GeminiGenerationConfig>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(rename = "systemInstruction")]
    system_instruction: Option<GeminiContent>,
}

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

#[derive(Debug, Serialize, Deserialize, Clone)]
struct GeminiPart {
    text: String,
}

#[derive(Debug, Serialize)]
struct GeminiGenerationConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(rename = "maxOutputTokens")]
    max_output_tokens: Option<usize>,
}

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

#[derive(Debug, Deserialize)]
struct GeminiCandidate {
    content: GeminiContent,
    #[allow(dead_code)]
    #[serde(rename = "finishReason")]
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
struct GeminiUsageMetadata {
    #[serde(rename = "promptTokenCount")]
    prompt_token_count: usize,
    #[serde(rename = "candidatesTokenCount")]
    candidates_token_count: usize,
    #[serde(rename = "totalTokenCount")]
    total_token_count: usize,
}

impl GeminiProvider {
    /// Create a new Gemini provider
    pub fn new() -> Result<Self> {
        let api_key = env::var("GEMINI_API_KEY")
            .or_else(|_| env::var("GOOGLE_API_KEY"))
            .map_err(|_| {
                Error::execution(
                    "GEMINI_API_KEY or GOOGLE_API_KEY environment variable not set",
                    "Set GEMINI_API_KEY or GOOGLE_API_KEY to use Gemini models",
                )
            })?;

        let client = reqwest::Client::new();

        Ok(Self { api_key, client })
    }

    /// Create a provider with explicit API key
    pub fn with_api_key(api_key: String) -> Self {
        let client = reqwest::Client::new();
        Self { api_key, client }
    }

    /// Convert generic messages to Gemini format
    fn convert_messages(&self, messages: Vec<crate::llm::LlmMessage>) -> Vec<GeminiContent> {
        messages
            .into_iter()
            .filter(|m| m.role != MessageRole::System)
            .map(|m| GeminiContent {
                role: match m.role {
                    MessageRole::User => "user".to_string(),
                    MessageRole::Assistant => "model".to_string(), // Gemini uses "model" not "assistant"
                    MessageRole::System => "user".to_string(),     // Fallback
                },
                parts: vec![GeminiPart { text: m.content }],
            })
            .collect()
    }

    /// Build API URL for the given model
    fn build_url(&self, model: &str, streaming: bool) -> String {
        let endpoint = if streaming {
            "streamGenerateContent"
        } else {
            "generateContent"
        };
        format!(
            "{}/{}:{}?key={}",
            GEMINI_API_BASE, model, endpoint, self.api_key
        )
    }
}

#[async_trait]
impl LlmProvider for GeminiProvider {
    async fn complete(&self, request: LlmRequest) -> Result<LlmResponse> {
        // Convert messages
        let contents = self.convert_messages(request.messages);

        // Build generation config
        let generation_config = if request.max_tokens.is_some() || request.temperature.is_some() {
            Some(GeminiGenerationConfig {
                temperature: request.temperature,
                max_output_tokens: request.max_tokens,
            })
        } else {
            None
        };

        // Build system instruction
        let system_instruction = request.system.map(|content| GeminiContent {
            role: "user".to_string(),
            parts: vec![GeminiPart { text: content }],
        });

        // Build Gemini request
        let gemini_request = GeminiRequest {
            contents,
            generation_config,
            system_instruction,
        };

        // Make API call
        let url = self.build_url(&request.model, false);
        let response = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&gemini_request)
            .send()
            .await
            .map_err(|e| Error::execution("Failed to call Gemini API", e.to_string()))?;

        // Check for HTTP errors
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(Error::execution(
                format!("Gemini API error: {}", status),
                error_text,
            ));
        }

        // Parse response
        let gemini_response: GeminiResponse = response
            .json()
            .await
            .map_err(|e| Error::execution("Failed to parse Gemini API response", e.to_string()))?;

        // Extract content from first candidate
        let content = gemini_response
            .candidates
            .first()
            .and_then(|candidate| {
                candidate
                    .content
                    .parts
                    .first()
                    .map(|part| part.text.clone())
            })
            .unwrap_or_default();

        // Extract usage metadata
        let usage = gemini_response.usage_metadata.map(|u| TokenUsage {
            input_tokens: u.prompt_token_count,
            output_tokens: u.candidates_token_count,
            total_tokens: u.total_token_count,
        });

        Ok(LlmResponse {
            content,
            model: request.model,
            usage,
        })
    }

    async fn stream(&self, request: LlmRequest) -> Result<LlmStream> {
        use futures::stream;

        // Convert messages
        let contents = self.convert_messages(request.messages.clone());

        // Build generation config
        let generation_config = if request.max_tokens.is_some() || request.temperature.is_some() {
            Some(GeminiGenerationConfig {
                temperature: request.temperature,
                max_output_tokens: request.max_tokens,
            })
        } else {
            None
        };

        // Build system instruction
        let system_instruction = request.system.map(|content| GeminiContent {
            role: "user".to_string(),
            parts: vec![GeminiPart { text: content }],
        });

        // Build Gemini streaming request
        let gemini_request = GeminiRequest {
            contents,
            generation_config,
            system_instruction,
        };

        // Make streaming API call
        let url = self.build_url(&request.model, true);
        let response = self
            .client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&gemini_request)
            .send()
            .await
            .map_err(|e| Error::execution("Failed to call Gemini API", e.to_string()))?;

        // Check for HTTP errors
        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(Error::execution(
                format!("Gemini API error: {}", status),
                error_text,
            ));
        }

        // Convert response to SSE stream
        use futures::TryStreamExt;
        let byte_stream = response.bytes_stream();
        let model = request.model.clone();

        let sse_stream = byte_stream
            .map_err(|e| Error::execution("Stream error", e.to_string()))
            .map(move |chunk_result| {
                match chunk_result {
                    Ok(bytes) => {
                        // Parse SSE events
                        let text = String::from_utf8_lossy(&bytes);
                        parse_sse_events(&text, &model)
                    }
                    Err(e) => vec![Err(e)],
                }
            })
            .flat_map(stream::iter);

        Ok(Box::pin(sse_stream))
    }

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

/// Parse SSE events from Gemini streaming response
fn parse_sse_events(text: &str, model: &str) -> Vec<Result<StreamChunk>> {
    let mut chunks = Vec::new();

    for line in text.lines() {
        // Gemini streams JSON objects separated by newlines (not SSE format with "data: " prefix)
        // Each line is a complete JSON response
        if line.trim().is_empty() || line.trim().starts_with('[') {
            continue;
        }

        // Parse JSON event
        if let Ok(event) = serde_json::from_str::<serde_json::Value>(line) {
            // Extract content from candidates
            if let Some(candidates) = event.get("candidates").and_then(|c| c.as_array()) {
                if let Some(candidate) = candidates.first() {
                    if let Some(content) = candidate.get("content") {
                        if let Some(parts) = content.get("parts").and_then(|p| p.as_array()) {
                            if let Some(part) = parts.first() {
                                if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
                                    chunks.push(Ok(StreamChunk::Content(text.to_string())));
                                }
                            }
                        }
                    }

                    // Check for finish reason (indicates stream end)
                    if let Some(_finish_reason) = candidate.get("finishReason") {
                        // Extract usage metadata if available
                        let usage = event.get("usageMetadata").and_then(|u| {
                            Some(TokenUsage {
                                input_tokens: u.get("promptTokenCount")?.as_u64()? as usize,
                                output_tokens: u.get("candidatesTokenCount")?.as_u64()? as usize,
                                total_tokens: u.get("totalTokenCount")?.as_u64()? as usize,
                            })
                        });

                        chunks.push(Ok(StreamChunk::Done(StreamMetadata {
                            model: model.to_string(),
                            usage,
                        })));
                    }
                }
            }

            // Check for errors
            if let Some(error) = event.get("error") {
                let error_msg = error
                    .get("message")
                    .and_then(|m| m.as_str())
                    .unwrap_or("Unknown error");
                chunks.push(Ok(StreamChunk::Error(error_msg.to_string())));
            }
        }
    }

    chunks
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::LlmMessage;

    #[test]
    fn test_create_provider_without_api_key() {
        // Should fail if GEMINI_API_KEY not set
        let original_gemini = env::var("GEMINI_API_KEY").ok();
        let original_google = env::var("GOOGLE_API_KEY").ok();
        env::remove_var("GEMINI_API_KEY");
        env::remove_var("GOOGLE_API_KEY");

        let result = GeminiProvider::new();
        assert!(result.is_err());

        // Restore original keys if they existed
        if let Some(key) = original_gemini {
            env::set_var("GEMINI_API_KEY", key);
        }
        if let Some(key) = original_google {
            env::set_var("GOOGLE_API_KEY", key);
        }
    }

    #[test]
    fn test_create_provider_with_explicit_key() {
        let provider = GeminiProvider::with_api_key("test-key".to_string());
        assert_eq!(provider.api_key, "test-key");
        assert_eq!(provider.name(), "Gemini");
    }

    #[test]
    fn test_convert_messages() {
        let provider = GeminiProvider::with_api_key("test-key".to_string());

        let messages = vec![
            LlmMessage {
                role: MessageRole::User,
                content: "Hello".to_string(),
            },
            LlmMessage {
                role: MessageRole::Assistant,
                content: "Hi there".to_string(),
            },
        ];

        let gemini_messages = provider.convert_messages(messages);

        assert_eq!(gemini_messages.len(), 2);
        assert_eq!(gemini_messages[0].role, "user");
        assert_eq!(gemini_messages[0].parts[0].text, "Hello");
        assert_eq!(gemini_messages[1].role, "model");
        assert_eq!(gemini_messages[1].parts[0].text, "Hi there");
    }

    #[test]
    fn test_parse_sse_content() {
        let json_line =
            r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"}}]}"#;
        let model = "gemini-2.0-flash-exp";

        let chunks = parse_sse_events(json_line, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Hello"),
            _ => panic!("Expected Content chunk"),
        }
    }

    #[test]
    fn test_parse_sse_with_finish_reason() {
        let json_line = r#"{"candidates":[{"content":{"parts":[{"text":"Done"}],"role":"model"},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"totalTokenCount":30}}"#;
        let model = "gemini-2.0-flash-exp";

        let chunks = parse_sse_events(json_line, model);

        assert_eq!(chunks.len(), 2);

        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Done"),
            _ => panic!("Expected Content chunk"),
        }

        match &chunks[1] {
            Ok(StreamChunk::Done(metadata)) => {
                assert_eq!(metadata.model, model);
                let usage = metadata.usage.as_ref().unwrap();
                assert_eq!(usage.input_tokens, 10);
                assert_eq!(usage.output_tokens, 20);
                assert_eq!(usage.total_tokens, 30);
            }
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_sse_error() {
        let json_line = r#"{"error":{"message":"API key invalid"}}"#;
        let model = "gemini-2.0-flash-exp";

        let chunks = parse_sse_events(json_line, model);

        assert_eq!(chunks.len(), 1);
        match &chunks[0] {
            Ok(StreamChunk::Error(msg)) => assert_eq!(msg, "API key invalid"),
            _ => panic!("Expected Error chunk"),
        }
    }

    #[test]
    fn test_parse_sse_multiple_chunks() {
        let json_lines = r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"}}]}
{"candidates":[{"content":{"parts":[{"text":" world"}],"role":"model"}}]}
{"candidates":[{"content":{"parts":[{"text":"!"}],"role":"model"},"finishReason":"STOP"}]}"#;
        let model = "gemini-2.0-flash-exp";

        let chunks = parse_sse_events(json_lines, model);

        assert_eq!(chunks.len(), 4);

        match &chunks[0] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "Hello"),
            _ => panic!("Expected first Content chunk"),
        }

        match &chunks[1] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, " world"),
            _ => panic!("Expected second Content chunk"),
        }

        match &chunks[2] {
            Ok(StreamChunk::Content(text)) => assert_eq!(text, "!"),
            _ => panic!("Expected third Content chunk"),
        }

        match &chunks[3] {
            Ok(StreamChunk::Done(_)) => {}
            _ => panic!("Expected Done chunk"),
        }
    }

    #[test]
    fn test_parse_sse_empty_input() {
        let json_lines = "";
        let model = "gemini-2.0-flash-exp";

        let chunks = parse_sse_events(json_lines, model);

        assert_eq!(chunks.len(), 0);
    }

    #[test]
    fn test_build_url_non_streaming() {
        let provider = GeminiProvider::with_api_key("test-key-123".to_string());
        let url = provider.build_url("gemini-2.0-flash-exp", false);

        assert!(url.contains("gemini-2.0-flash-exp:generateContent"));
        assert!(url.contains("key=test-key-123"));
    }

    #[test]
    fn test_build_url_streaming() {
        let provider = GeminiProvider::with_api_key("test-key-456".to_string());
        let url = provider.build_url("gemini-2.0-flash-exp", true);

        assert!(url.contains("gemini-2.0-flash-exp:streamGenerateContent"));
        assert!(url.contains("key=test-key-456"));
    }

    #[tokio::test]
    #[ignore] // Only run with real API key
    async fn test_gemini_api_call() {
        let provider = GeminiProvider::new().expect("GEMINI_API_KEY must be set");

        let request = LlmRequest {
            model: "gemini-2.0-flash-exp".to_string(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: "Say hello in exactly 3 words.".to_string(),
            }],
            max_tokens: Some(50),
            temperature: Some(0.7),
            system: None,
        };

        let response = provider.complete(request).await;
        assert!(response.is_ok());

        let response = response.unwrap();
        assert!(!response.content.is_empty());
        assert!(response.usage.is_some());

        println!("Response: {}", response.content);
        println!("Usage: {:?}", response.usage);
    }
}
