//! Google Gemini provider
//!
//! This provider supports Google's Gemini models via the Generative Language API.
//! It can be configured either programmatically or via config.yml.

use crate::config::ModelConfig;
use crate::http::{post_json, HttpRequestSpec};
use async_trait::async_trait;
use enact_core::providers::{
    ChatChoice, ChatMessage, ChatRequest, ChatResponse, ChatUsage, ModelCapabilities, ModelProvider,
};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;

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

#[derive(Debug, Serialize)]
struct GeminiContent {
    #[serde(skip_serializing_if = "Option::is_none")]
    role: Option<String>,
    parts: Vec<GeminiPart>,
}

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

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

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

#[derive(Debug, Deserialize)]
struct GeminiResponse {
    #[serde(default)]
    candidates: Vec<GeminiCandidate>,
    #[serde(default)]
    usage_metadata: Option<GeminiUsageMetadata>,
}

#[derive(Debug, Deserialize)]
struct GeminiCandidate {
    content: GeminiCandidateContent,
}

#[derive(Debug, Deserialize)]
struct GeminiCandidateContent {
    #[serde(default)]
    parts: Vec<GeminiPartIn>,
}

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

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

/// Google Gemini provider.
///
/// Supports Gemini models via the Generative Language API.
/// All configuration can be loaded from config.yml via `from_model_config()`.
pub struct Gemini {
    client: Client,
    /// Model identifier
    model_id: String,
    /// Full URL for API calls (includes auth when required)
    full_url: String,
    /// Request timeout
    timeout: Duration,
    /// Whether temperature is supported
    has_temperature_supported: bool,
    /// Default max tokens
    max_tokens_default: u32,
    /// Default temperature
    temperature_default: f32,
    /// Whether tool calls are supported
    supports_tool_calls: bool,
    /// Cost per 1K input tokens (from config cost.cost_per_1m_input)
    cost_per_1m_input: Option<f64>,
    /// Cost per 1K output tokens (from config cost.cost_per_1m_output)
    cost_per_1m_output: Option<f64>,
    /// Cost per 1K pixels for image (from config cost.cost_per_1m_pixels)
    cost_per_1m_pixels: Option<f64>,
}

impl Gemini {
    /// Create a new Gemini provider from model configuration.
    ///
    /// # Arguments
    ///
    /// * `model_id` - The model identifier
    /// * `config` - The model configuration from config.yml
    /// * `api_key` - The resolved API key
    pub fn from_model_config(
        model_id: &str,
        config: &ModelConfig,
        api_key: &str,
    ) -> anyhow::Result<Self> {
        // Build the full URL with API key
        let full_url = if config.runtime.full_url.contains("key=") {
            config.runtime.full_url.clone()
        } else {
            format!("{}?key={}", config.runtime.full_url, api_key)
        };

        Ok(Self {
            client: Client::new(),
            model_id: model_id.to_string(),
            full_url,
            timeout: Duration::from_secs(config.runtime.timeout_seconds),
            has_temperature_supported: config.runtime.has_temperature_supported,
            max_tokens_default: config.runtime.max_tokens_default,
            temperature_default: config.runtime.temperature_default,
            supports_tool_calls: config.capabilities.tool_calls,
            cost_per_1m_input: config.cost.cost_per_1m_input,
            cost_per_1m_output: config.cost.cost_per_1m_output,
            cost_per_1m_pixels: config.cost.cost_per_1m_pixels,
        })
    }

    #[allow(dead_code)] // used in tests
    fn format_model(model: &str) -> String {
        if model.starts_with("models/") {
            model.to_string()
        } else {
            format!("models/{model}")
        }
    }

    fn convert_messages(
        messages: &[ChatMessage],
    ) -> (Option<GeminiSystemInstruction>, Vec<GeminiContent>) {
        let mut system = None;
        let mut contents = Vec::new();

        for msg in messages {
            if msg.role == "system" && system.is_none() {
                system = Some(GeminiSystemInstruction {
                    parts: vec![GeminiPart {
                        text: msg.content.clone().unwrap_or_default(),
                    }],
                });
                continue;
            }

            let role = if msg.role == "assistant" {
                "model"
            } else {
                "user"
            };
            contents.push(GeminiContent {
                role: Some(role.to_string()),
                parts: vec![GeminiPart {
                    text: msg.content.clone().unwrap_or_default(),
                }],
            });
        }

        (system, contents)
    }
}

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

    fn model(&self) -> &str {
        &self.model_id
    }

    fn capabilities(&self) -> ModelCapabilities {
        ModelCapabilities {
            max_tokens: self.max_tokens_default,
            max_output_tokens: self.max_tokens_default,
            supports_streaming: false,
            supports_tools: self.supports_tool_calls,
            supports_reasoning: false,
            supports_vision: true,
            supports_json_mode: true,
            supports_embeddings: false,
            supports_image_generation: false,
            supports_audio_transcription: false,
            supports_speech: false,
            supports_video_generation: false,
            pii_safe: false,
            cost_per_1m_input: self.cost_per_1m_input,
            cost_per_1m_output: self.cost_per_1m_output,
            cost_per_1m_pixels: self.cost_per_1m_pixels,
        }
    }

    async fn chat(&self, request: ChatRequest) -> anyhow::Result<ChatResponse> {
        let (system_instruction, contents) = Self::convert_messages(&request.messages);

        // Build temperature (only include if supported)
        let temperature = if self.has_temperature_supported {
            Some(request.temperature.unwrap_or(self.temperature_default))
        } else {
            None
        };

        let body = GeminiRequest {
            contents,
            system_instruction,
            generation_config: GeminiGenerationConfig {
                temperature,
                max_output_tokens: request.max_tokens.unwrap_or(self.max_tokens_default),
            },
        };

        let spec = HttpRequestSpec {
            url: self.full_url.clone(),
            headers: vec![("content-type".to_string(), "application/json".to_string())],
            body: serde_json::to_value(body)?,
            timeout: self.timeout,
        };

        let response: GeminiResponse =
            serde_json::from_value(post_json(&self.client, spec).await?)?;
        let text = response
            .candidates
            .into_iter()
            .next()
            .map(|candidate| {
                candidate
                    .content
                    .parts
                    .into_iter()
                    .filter_map(|part| part.text)
                    .collect::<Vec<_>>()
                    .join("\n")
            })
            .unwrap_or_default();

        let usage = response.usage_metadata.map(|u| ChatUsage {
            prompt_tokens: u.prompt_token_count,
            completion_tokens: u.candidates_token_count,
            total_tokens: u.total_token_count,
        });

        Ok(ChatResponse {
            id: format!("gemini-{}", uuid::Uuid::new_v4()),
            choices: vec![ChatChoice {
                index: 0,
                message: ChatMessage::assistant(text),
                finish_reason: Some("stop".to_string()),
            }],
            usage,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{
        AuthConfig, AuthType, CapabilitiesConfig, CostConfig, GovernanceConfig, LifecycleStatus,
        Modality, Protocol, RuntimeConfig, SecretType, SecretsConfig,
    };

    fn create_test_config() -> ModelConfig {
        ModelConfig {
            display_name: "Gemini Test".to_string(),
            provider_name: "gemini".to_string(),
            input_modality: vec![Modality::Text],
            output_modality: vec![Modality::Text],
            runtime: RuntimeConfig {
                base_url: "https://generativelanguage.googleapis.com/v1beta".to_string(),
                full_url: "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent".to_string(),
                protocol: Protocol::OpenAICompatible,
                api_version: "v1beta".to_string(),
                use_max_completion_tokens: false,
                has_temperature_supported: true,
                additional_notes: vec![],
                model_param: false,
                timeout_seconds: 60,
                max_tokens_default: 4096,
                temperature_default: 0.7,
            },
            auth: AuthConfig {
                auth_type: AuthType::Bearer,
            },
            secrets: SecretsConfig {
                secret_type: SecretType::Environment,
                name: "GOOGLE_API_KEY".to_string(),
            },
            capabilities: CapabilitiesConfig {
                tool_calls: true,
                reasoning: false,
            },
            governance: GovernanceConfig {
                is_use_approve: true,
                approved_by: Some("admin".to_string()),
                approval_date: Some("2024-01-01".to_string()),
                lifecycle_status: LifecycleStatus::Active,
            },
            cost: CostConfig::default(),
            fallbacks: vec![],
        }
    }

    #[test]
    fn test_from_model_config() {
        let config = create_test_config();
        let provider =
            Gemini::from_model_config("gemini:text:gemini-pro", &config, "test-key").unwrap();

        assert_eq!(provider.name(), "gemini");
        assert_eq!(provider.model(), "gemini:text:gemini-pro");
        assert!(provider.full_url.contains("key=test-key"));
    }

    #[test]
    fn test_format_model() {
        assert_eq!(Gemini::format_model("gemini-pro"), "models/gemini-pro");
        assert_eq!(
            Gemini::format_model("models/gemini-pro"),
            "models/gemini-pro"
        );
    }

    #[test]
    fn test_convert_messages_extracts_system() {
        let messages = vec![
            ChatMessage::system("You are a helpful assistant"),
            ChatMessage::user("Hello"),
        ];

        let (system, contents) = Gemini::convert_messages(&messages);

        assert!(system.is_some());
        assert_eq!(contents.len(), 1);
        assert_eq!(contents[0].role, Some("user".to_string()));
    }

    #[test]
    fn test_capabilities_from_config() {
        let config = create_test_config();
        let provider = Gemini::from_model_config("test", &config, "key").unwrap();

        let caps = provider.capabilities();
        assert!(caps.supports_tools);
        assert!(caps.supports_vision);
    }
}
