//! Google Gemini API client implementation for chat and completion functionality.
//!
//! This module provides integration with Google's Gemini models through their API.
//! It implements chat, completion and embedding capabilities via the Gemini API.
//!
//! # Features
//! - Chat conversations with system prompts and message history
//! - Text completion requests
//! - Configuration options for temperature, tokens, top_p, top_k etc.
//! - Streaming support
//!
//! # Example
//! ```no_run
//! use llm::backends::google::Google;
//! use llm::chat::{ChatMessage, ChatRole, ChatProvider, Tool};
//!
//! #[tokio::main]
//! async fn main() {
//! let client = Google::new(
//!     "your-api-key",
//!     None, // Use default model
//!     Some(1000), // Max tokens
//!     Some(0.7), // Temperature
//!     None, // Default timeout
//!     None, // No system prompt
//!     None, // No streaming
//!     None, // Default top_p
//!     None, // Default top_k
//!     None, // No JSON schema
//!     None, // No tools
//! );
//!
//! let messages = vec![
//!     ChatMessage::user().content("Hello!").build()
//! ];
//!
//! let response = client.chat(&messages).await.unwrap();
//! println!("{}", response);
//! }
//! ```

use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use http::{Method, Request, Response, header::CONTENT_TYPE};
use querymt::{
    FunctionCall, HTTPLLMProvider, ToolCall, Usage,
    auth::ApiKeyResolver,
    chat::{
        ChatMessage, ChatResponse, ChatRole, Content, FinishReason, ReasoningEffort, StreamChunk,
        StructuredOutputFormat, Tool, ToolChoice,
        http::{ChatStreamParser, HTTPChatProvider},
    },
    completion::{CompletionRequest, CompletionResponse, http::HTTPCompletionProvider},
    embedding::http::HTTPEmbeddingProvider,
    error::LLMError,
    handle_http_error,
    plugin::HTTPLLMProviderFactory,
};
use schemars::{JsonSchema, schema_for};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
use url::Url;

/// Client for interacting with Google's Gemini API.
///
/// This struct holds the configuration and state needed to make requests to the Gemini API.
/// It implements the [`ChatProvider`], [`CompletionProvider`], and [`EmbeddingProvider`] traits.
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub struct Google {
    /// API key for authentication with Google's API
    pub api_key: String,
    /// Model identifier (e.g. "gemini-1.5-flash")
    pub model: String,
    /// Maximum number of tokens to generate in responses
    pub max_tokens: Option<u32>,
    /// Sampling temperature between 0.0 and 1.0
    pub temperature: Option<f32>,
    /// Optional system prompt to set context
    #[serde(
        default,
        deserialize_with = "querymt::params::deserialize_system_string"
    )]
    pub system: Option<String>,
    /// Request timeout in seconds
    pub timeout_seconds: Option<u64>,
    /// Whether to stream responses
    pub stream: Option<bool>,
    /// Top-p sampling parameter
    pub top_p: Option<f32>,
    /// Top-k sampling parameter
    pub top_k: Option<u32>,
    /// JSON schema for structured output
    pub json_schema: Option<StructuredOutputFormat>,
    /// Available tools for function calling
    pub tools: Option<Vec<Tool>>,
    pub tool_choice: Option<ToolChoice>, // FIXME: currently not being used
    pub reasoning_effort: Option<ReasoningEffort>,
    pub thinking_budget: Option<u32>,
    pub cached_content: Option<String>,

    /// Optional resolver for dynamic credential refresh (e.g., OAuth tokens).
    #[serde(skip)]
    #[schemars(skip)]
    pub key_resolver: Option<Arc<dyn ApiKeyResolver>>,
}

#[derive(Serialize)]
struct GoogleChatRequest<'a> {
    /// List of conversation messages
    contents: Vec<GoogleChatContent<'a>>,
    /// Optional generation parameters
    #[serde(skip_serializing_if = "Option::is_none", rename = "generationConfig")]
    generation_config: Option<GoogleGenerationConfig>,
    /// Tools that the model can use
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<GoogleTool>>,
    /// Tools choice
    #[serde(skip_serializing_if = "Option::is_none")]
    tool_config: Option<GoogleToolConfig>,
    /// Tools choice
    #[serde(skip_serializing_if = "Option::is_none", rename = "cachedContent")]
    cached_content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system_instruction: Option<GoogleSystemInstruction<'a>>,
}

/// Individual message in a chat conversation
#[derive(Serialize)]
struct GoogleChatContent<'a> {
    /// Role of the message sender ("user", "model", or "system")
    role: &'a str,
    /// Content parts of the message
    parts: Vec<GoogleContentPart<'a>>,
}

#[derive(Serialize)]
struct GoogleSystemInstruction<'a> {
    /// Content parts of the message
    parts: Vec<GoogleContentPart<'a>>,
}

/// Individual part of a content message
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct GoogleContentPart<'a> {
    /// The actual text content
    #[serde(skip_serializing_if = "Option::is_none")]
    text: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    inline_data: Option<GoogleInlineData>,
    #[serde(skip_serializing_if = "Option::is_none")]
    function_call: Option<GoogleFunctionCall>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "functionResponse")]
    function_response: Option<GoogleFunctionResponse>,
    #[serde(skip_serializing_if = "Option::is_none")]
    thought: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "thoughtSignature")]
    thought_signature: Option<String>,
}

impl<'a> GoogleContentPart<'a> {
    fn text(text: &'a str) -> Self {
        Self {
            text: Some(text),
            inline_data: None,
            function_call: None,
            function_response: None,
            thought: None,
            thought_signature: None,
        }
    }

    fn thought(text: &'a str) -> Self {
        Self {
            text: Some(text),
            inline_data: None,
            function_call: None,
            function_response: None,
            thought: Some(true),
            thought_signature: None,
        }
    }

    fn inline_data(mime_type: String, data: String) -> Self {
        Self {
            text: None,
            inline_data: Some(GoogleInlineData { mime_type, data }),
            function_call: None,
            function_response: None,
            thought: None,
            thought_signature: None,
        }
    }

    fn function_call(name: String, args: Value, thought_signature: Option<String>) -> Self {
        Self {
            text: None,
            inline_data: None,
            function_call: Some(GoogleFunctionCall { name, args }),
            function_response: None,
            thought: None,
            thought_signature,
        }
    }

    fn function_response(name: String, content: Value) -> Self {
        Self {
            text: None,
            inline_data: None,
            function_call: None,
            function_response: Some(GoogleFunctionResponse {
                name: name.clone(),
                response: GoogleFunctionResponseContent { name, content },
            }),
            thought: None,
            thought_signature: None,
        }
    }
}

#[derive(Serialize)]
struct GoogleInlineData {
    mime_type: String,
    data: String,
}

/// Configuration parameters for text generation
#[derive(Serialize)]
struct GoogleGenerationConfig {
    /// Maximum tokens to generate
    #[serde(skip_serializing_if = "Option::is_none", rename = "maxOutputTokens")]
    max_output_tokens: Option<u32>,
    /// Sampling temperature
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    /// Top-p sampling parameter
    #[serde(skip_serializing_if = "Option::is_none", rename = "topP")]
    top_p: Option<f32>,
    /// Top-k sampling parameter
    #[serde(skip_serializing_if = "Option::is_none", rename = "topK")]
    top_k: Option<u32>,
    /// The MIME type of the response
    #[serde(skip_serializing_if = "Option::is_none")]
    response_mime_type: Option<GoogleResponseMimeType>,
    /// A schema for structured output
    #[serde(skip_serializing_if = "Option::is_none")]
    response_schema: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "thinkingConfig")]
    thinking_config: Option<GoogleThinkingConfig>,
}

/// Configuration parameters for text generation
#[derive(Serialize)]
struct GoogleThinkingConfig {
    #[serde(skip_serializing_if = "Option::is_none", rename = "includeThoughts")]
    include_thoughts: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "thinkingBudget")]
    thinking_budget: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "thinkingLevel")]
    thinking_level: Option<String>,
}

/// Response from the chat completion API
#[derive(Deserialize, Debug)]
struct GoogleChatResponse {
    /// Generated completion candidates
    candidates: Vec<GoogleCandidate>,
    #[serde(rename = "usageMetadata")]
    usage: Option<Usage>,
}

impl std::fmt::Display for GoogleChatResponse {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match (self.text(), self.tool_calls()) {
            (Some(text), Some(tool_calls)) => {
                for call in tool_calls {
                    write!(f, "{:?}", call)?;
                }
                write!(f, "{}", text)
            }
            (Some(text), None) => write!(f, "{}", text),
            (None, Some(tool_calls)) => {
                for call in tool_calls {
                    write!(f, "{:?}", call)?;
                }
                Ok(())
            }
            (None, None) => write!(f, ""),
        }
    }
}

/// Individual completion candidate
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
struct GoogleCandidate {
    /// Content of the candidate response
    content: GoogleResponseContent,
    /// Finish reason (only present in final streaming chunk or complete response)
    finish_reason: Option<String>,
    /// Index of this candidate
    index: usize,
}

/// Response content block
#[derive(Deserialize, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct GoogleResponseContent {
    /// Parts making up the content
    #[serde(default)]
    parts: Vec<GoogleResponsePart>,
    /// Function calls if any are used
    #[serde(skip_serializing_if = "Option::is_none")]
    function_call: Option<GoogleFunctionCall>,
    /// Function calls as array
    #[serde(skip_serializing_if = "Option::is_none")]
    function_calls: Option<Vec<GoogleFunctionCall>>,
}

impl ChatResponse for GoogleChatResponse {
    fn text(&self) -> Option<String> {
        self.candidates.first().map(|c| {
            c.content
                .parts
                .iter()
                .filter(|p| !p.thought)
                .map(|p| p.text.clone().unwrap_or_default())
                .collect()
        })
    }

    fn thinking(&self) -> Option<String> {
        self.candidates.first().and_then(|c| {
            let thoughts: String = c
                .content
                .parts
                .iter()
                .filter(|p| p.thought)
                .map(|p| p.text.clone().unwrap_or_default())
                .collect();
            if thoughts.is_empty() {
                None
            } else {
                Some(thoughts)
            }
        })
    }

    fn tool_calls(&self) -> Option<Vec<ToolCall>> {
        self.candidates.first().and_then(|c| {
            // First check for function calls at the part level (new API format)
            let part_function_calls: Vec<ToolCall> = c
                .content
                .parts
                .iter()
                .filter_map(|part| {
                    part.function_call.as_ref().map(|f| {
                        let id = if let Some(sig) = &part.thought_signature {
                            format!("call_{}:{}", f.name, sig)
                        } else {
                            format!("call_{}", f.name)
                        };

                        ToolCall {
                            id,
                            call_type: "function".to_string(),
                            function: FunctionCall {
                                name: f.name.clone(),
                                arguments: serde_json::to_string(&f.args).unwrap_or_default(),
                            },
                        }
                    })
                })
                .collect();

            if !part_function_calls.is_empty() {
                return Some(part_function_calls);
            }

            // Otherwise check for function_calls/function_call at the content level (older format)
            if let Some(fc) = &c.content.function_calls {
                // Process array of function calls
                Some(
                    fc.iter()
                        .map(|f| {
                            let id = format!("call_{}", f.name);

                            ToolCall {
                                id,
                                call_type: "function".to_string(),
                                function: FunctionCall {
                                    name: f.name.clone(),
                                    arguments: serde_json::to_string(&f.args).unwrap_or_default(),
                                },
                            }
                        })
                        .collect(),
                )
            } else {
                c.content.function_call.as_ref().map(|f| {
                    let id = format!("call_{}", f.name);

                    vec![ToolCall {
                        id,
                        call_type: "function".to_string(),
                        function: FunctionCall {
                            name: f.name.clone(),
                            arguments: serde_json::to_string(&f.args).unwrap_or_default(),
                        },
                    }]
                })
            }
        })
    }

    fn usage(&self) -> Option<Usage> {
        self.usage.clone()
    }

    fn finish_reason(&self) -> Option<FinishReason> {
        if self.tool_calls().is_some() {
            return Some(FinishReason::ToolCalls);
        }

        match self
            .candidates
            .first()
            .map(|c| c.finish_reason.clone())
            .unwrap()
            .as_deref()
        {
            Some("STOP") => Some(FinishReason::Stop),
            Some("MAX_TOKENS") => Some(FinishReason::Length),
            Some(
                "SAFETY"
                | "IMAGE_SAFETY"
                | "IMAGE_PROHIBITED_CONTENT"
                | "BLOCKLIST"
                | "SPII"
                | "PROHIBITED_CONTENT",
            ) => Some(FinishReason::ContentFilter),
            Some("OTHER" | "IMAGE_OTHER") => Some(FinishReason::Other),
            Some("MALFORMED_FUNCTION_CALL") => Some(FinishReason::Error),
            Some("FINISH_REASON_UNSPECIFIED") => Some(FinishReason::Unknown),
            _ => None,
        }
    }
}

/// Individual part of response content
#[derive(Deserialize, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct GoogleResponsePart {
    /// Text content of this part (may be absent if functionCall is present)
    text: Option<String>,
    /// Function call contained in this part
    function_call: Option<GoogleFunctionCall>,
    /// Whether this part is a thought part
    #[serde(default)]
    thought: bool,
    /// Thought signature for the function call
    #[serde(skip_serializing_if = "Option::is_none", rename = "thoughtSignature")]
    thought_signature: Option<String>,
}

/// Google's function calling tool definition
#[derive(Serialize, Debug)]
struct GoogleTool {
    /// The function declarations array
    #[serde(rename = "functionDeclarations")]
    function_declarations: Vec<GoogleFunctionDeclaration>,
}

/// MIME type of the response
#[derive(Deserialize, Debug, Serialize)]
enum GoogleResponseMimeType {
    /// Plain text response
    #[serde(rename = "text/plain")]
    PlainText,
    /// JSON response
    #[serde(rename = "application/json")]
    Json,
    /// ENUM as a string response in the response candidates.
    #[serde(rename = "text/x.enum")]
    Enum,
}

/// Google's tool configuration
#[derive(Serialize, Debug)]
struct GoogleToolConfig {
    #[serde(
        skip_serializing_if = "Option::is_none",
        rename = "functionCallingConfig"
    )]
    function_calling_config: Option<FunctionCallingConfig>,
}

#[derive(Serialize, Debug)]
struct FunctionCallingConfig {
    #[serde(skip_serializing_if = "Option::is_none")]
    mode: Option<String>,
    #[serde(
        skip_serializing_if = "Option::is_none",
        rename = "allowedFunctionNames"
    )]
    allowed_function_names: Option<Vec<String>>,
}

/// Google function declaration, similar to OpenAI's function definition
#[derive(Serialize, Debug)]
struct GoogleFunctionDeclaration {
    /// Name of the function
    name: String,
    /// Description of what the function does
    description: String,
    /// Parameters for the function
    parameters: GoogleFunctionParameters,
}

impl From<&Tool> for GoogleFunctionDeclaration {
    fn from(tool: &Tool) -> Self {
        fn remove_additional(v: &mut Value) {
            match v {
                Value::Object(map) => {
                    // Remove this object's `additionalProperties` key, if present.
                    map.remove("additionalProperties");
                    // Recurse into all remaining values
                    for child in map.values_mut() {
                        remove_additional(child);
                    }
                }
                Value::Array(arr) => {
                    // If it's an array, recurse into each element
                    for child in arr.iter_mut() {
                        remove_additional(child);
                    }
                }
                _ => {}
            }
        }

        let properties_value = tool
            .function
            .parameters
            .get("properties")
            .map(|p| {
                let mut props = p.clone();
                remove_additional(&mut props);
                props
            })
            .unwrap_or_else(|| serde_json::Value::Object(serde_json::Map::new()));

        GoogleFunctionDeclaration {
            name: tool.function.name.clone(),
            description: tool.function.description.clone(),
            parameters: GoogleFunctionParameters {
                schema_type: "object".to_string(),
                properties: properties_value,
                required: tool
                    .function
                    .parameters
                    .get("required")
                    .and_then(|v| v.as_array())
                    .map(|arr| {
                        arr.iter()
                            .filter_map(|v| v.as_str().map(|s| s.to_string()))
                            .collect::<Vec<String>>()
                    })
                    .unwrap_or_default(),
            },
        }
    }
}

/// Google function parameters schema
#[derive(Serialize, Debug)]
struct GoogleFunctionParameters {
    /// The type of parameters object (usually "object")
    #[serde(rename = "type")]
    schema_type: String,
    /// Map of parameter names to their properties
    properties: Value,
    /// List of required parameter names
    required: Vec<String>,
}

/// Google function call object in response
#[derive(Deserialize, Debug, Serialize)]
struct GoogleFunctionCall {
    /// Name of the function to call
    name: String,
    /// Arguments for the function call as structured JSON
    #[serde(default)]
    args: Value,
}

/// Google function response wrapper for function results
///
/// Format follows Google's Gemini API specification for function calling results:
/// https://ai.google.dev/docs/function_calling
///
/// The expected format is:
/// {
///   "role": "function",
///   "parts": [{
///     "functionResponse": {
///       "name": "function_name",
///       "response": {
///         "name": "function_name",
///         "content": { ... } // JSON content returned by the function
///       }
///     }
///   }]
/// }
#[derive(Deserialize, Debug, Serialize)]
struct GoogleFunctionResponse {
    /// Name of the function that was called
    name: String,
    /// Response from the function as structured JSON
    response: GoogleFunctionResponseContent,
}

#[derive(Deserialize, Debug, Serialize)]
struct GoogleFunctionResponseContent {
    /// Name of the function that was called
    name: String,
    /// Content of the function response
    content: Value,
}

/// Request body for embedding content
#[derive(Serialize)]
struct GoogleEmbeddingRequest<'a> {
    model: &'a str,
    content: GoogleEmbeddingContent<'a>,
}

#[derive(Serialize)]
struct GoogleEmbeddingContent<'a> {
    parts: Vec<GoogleContentPart<'a>>,
}

/// Response from the embedding API
#[derive(Deserialize)]
struct GoogleEmbeddingResponse {
    embedding: GoogleEmbedding,
}

#[derive(Deserialize)]
struct GoogleEmbedding {
    values: Vec<f32>,
}

impl Google {
    fn default_base_url() -> Url {
        Url::parse("https://generativelanguage.googleapis.com/v1beta/models/").unwrap()
    }

    /// Returns the current API key, using the resolver if available.
    fn resolved_key(&self) -> String {
        if let Some(ref resolver) = self.key_resolver {
            resolver.current()
        } else {
            self.api_key.clone()
        }
    }

    fn is_gemini_2_5(&self) -> bool {
        self.model.contains("2.5")
    }

    fn effort_to_budget(effort: ReasoningEffort) -> u32 {
        match effort {
            ReasoningEffort::Low => 1_024,
            ReasoningEffort::Medium => 8_000,
            ReasoningEffort::High => 16_000,
            ReasoningEffort::Max => 24_576,
        }
    }

    fn effort_to_level(effort: ReasoningEffort) -> String {
        match effort {
            ReasoningEffort::Low => "low".to_string(),
            ReasoningEffort::Medium | ReasoningEffort::High | ReasoningEffort::Max => {
                "high".to_string()
            }
        }
    }

    fn resolve_thinking_config(&self) -> Option<GoogleThinkingConfig> {
        if let Some(thinking_budget) = self.thinking_budget {
            return Some(GoogleThinkingConfig {
                include_thoughts: Some(true),
                thinking_budget: Some(thinking_budget),
                thinking_level: None,
            });
        }

        let effort = self.reasoning_effort?;
        if self.is_gemini_2_5() {
            Some(GoogleThinkingConfig {
                include_thoughts: Some(true),
                thinking_budget: Some(Self::effort_to_budget(effort)),
                thinking_level: None,
            })
        } else {
            Some(GoogleThinkingConfig {
                include_thoughts: Some(true),
                thinking_budget: None,
                thinking_level: Some(Self::effort_to_level(effort)),
            })
        }
    }
}

impl HTTPChatProvider for Google {
    fn chat_request(
        &self,
        messages: &[ChatMessage],
        tools: Option<&[Tool]>,
    ) -> Result<Request<Vec<u8>>, LLMError> {
        let resolved_key = self.resolved_key();
        if resolved_key.is_empty() {
            return Err(LLMError::AuthError("Missing Google API key".into()));
        }

        let mut chat_contents = Vec::with_capacity(messages.len());

        if self.cached_content.is_some() && self.tools().is_some() {
            return Err(LLMError::InvalidRequest(
                "Cached content is set and tools as well!".to_string(),
            ));
        }

        for msg in messages {
            let has_tool_result = msg.content.iter().any(|b| b.is_tool_result());
            let role = if has_tool_result {
                "function"
            } else {
                match msg.role {
                    ChatRole::User => "user",
                    ChatRole::Assistant => "model",
                }
            };

            let mut parts = Vec::new();

            for block in &msg.content {
                match block {
                    Content::Text { text } => {
                        if !text.is_empty() {
                            parts.push(GoogleContentPart::text(text));
                        }
                    }
                    Content::Thinking { text, .. } => {
                        if !text.is_empty() {
                            parts.push(GoogleContentPart::thought(text));
                        }
                    }
                    Content::Image { mime_type, data } => {
                        parts.push(GoogleContentPart::inline_data(
                            mime_type.clone(),
                            BASE64.encode(data),
                        ));
                    }
                    Content::Pdf { data } => {
                        parts.push(GoogleContentPart::inline_data(
                            "application/pdf".to_string(),
                            BASE64.encode(data),
                        ));
                    }
                    Content::ImageUrl { url } => {
                        // Google input parts do not expose a direct image URL field,
                        // so preserve the reference as text.
                        parts.push(GoogleContentPart::text(url));
                    }
                    Content::ToolUse {
                        id,
                        name,
                        arguments,
                    } => {
                        let expected_prefix = format!("call_{}:", name);
                        let signature = if id.starts_with(&expected_prefix) {
                            Some(id[expected_prefix.len()..].to_string())
                        } else {
                            None
                        };
                        parts.push(GoogleContentPart::function_call(
                            name.clone(),
                            arguments.clone(),
                            signature,
                        ));
                    }
                    Content::ToolResult {
                        id, name, content, ..
                    } => {
                        let text = content
                            .iter()
                            .filter_map(|c| c.as_text())
                            .collect::<Vec<_>>()
                            .join("\n");
                        let payload = if text.is_empty() {
                            Value::Null
                        } else {
                            Value::String(text)
                        };
                        parts.push(GoogleContentPart::function_response(
                            name.clone().unwrap_or_else(|| id.clone()),
                            payload,
                        ));
                    }
                    Content::Audio { .. } | Content::ResourceLink { .. } => {
                        // Unsupported in Google request format today.
                    }
                }
            }

            chat_contents.push(GoogleChatContent { role, parts });
        }

        // Add system message if present
        let system_instruction = self.system.as_ref().map(|system| GoogleSystemInstruction {
            parts: vec![GoogleContentPart::text(system)],
        });

        // Convert tools to Google's format if provided
        let google_tools = tools.map(|t| {
            vec![GoogleTool {
                function_declarations: t.iter().map(GoogleFunctionDeclaration::from).collect(),
            }]
        });

        // Build generation config
        let generation_config = {
            // If json_schema and json_schema.schema are not None, use json_schema.schema as the response schema and set response_mime_type to JSON
            // Google's API doesn't need the schema to have a "name" field, so we can just use the schema directly.
            let (response_mime_type, response_schema) = if let Some(json_schema) = &self.json_schema
            {
                if let Some(schema) = &json_schema.schema {
                    // If the schema has an "additionalProperties" field (as required by OpenAI), remove it as Google's API doesn't support it
                    let mut schema = schema.clone();

                    if let Some(obj) = schema.as_object_mut() {
                        obj.remove("additionalProperties");
                    }

                    (Some(GoogleResponseMimeType::Json), Some(schema))
                } else {
                    (None, None)
                }
            } else {
                (None, None)
            };

            let thinking_config = self.resolve_thinking_config();

            Some(GoogleGenerationConfig {
                max_output_tokens: self.max_tokens,
                temperature: self.temperature,
                top_p: self.top_p,
                top_k: self.top_k,
                response_mime_type,
                response_schema,
                thinking_config,
            })
        };

        let req_body = GoogleChatRequest {
            contents: chat_contents,
            generation_config,
            tools: google_tools,
            tool_config: None, // FIXME
            cached_content: self.cached_content.clone(),
            system_instruction,
        };

        let json_body = serde_json::to_vec(&req_body)?;

        // Use streamGenerateContent endpoint when streaming is enabled
        let endpoint = if self.stream.unwrap_or(false) {
            "streamGenerateContent"
        } else {
            "generateContent"
        };

        let path = format!(
            "{}{}:{}",
            Google::default_base_url().path(),
            &self.model,
            endpoint
        );
        let mut url = Google::default_base_url()
            .join(&path)
            .map_err(|e| LLMError::HttpError(e.to_string()))?;
        url.set_query(Some(&format!("key={}", &resolved_key)));

        Ok(Request::builder()
            .method(Method::POST)
            .uri(url.as_str())
            .header(CONTENT_TYPE, "application/json")
            .body(json_body)?)
    }

    fn chat_stream_request(
        &self,
        messages: &[ChatMessage],
        tools: Option<&[Tool]>,
    ) -> Result<Request<Vec<u8>>, LLMError> {
        let mut cfg = self.clone();
        cfg.stream = Some(true);
        cfg.chat_request(messages, tools)
    }

    fn parse_chat(&self, resp: Response<Vec<u8>>) -> Result<Box<dyn ChatResponse>, LLMError> {
        handle_http_error!(resp);

        let json_resp: Result<GoogleChatResponse, serde_json::Error> =
            serde_json::from_slice(resp.body());

        match json_resp {
            Ok(response) => Ok(Box::new(response)),
            Err(e) => {
                // Return a more descriptive error with the raw response
                Err(LLMError::ResponseFormatError {
                    message: format!("Failed to decode Google API response: {}", e),
                    raw_response: String::from_utf8(resp.into_body())?,
                })
            }
        }
    }

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

    fn chat_stream_parser(&self) -> Result<Box<dyn ChatStreamParser>, LLMError> {
        Ok(Box::new(GoogleStreamParser::default()))
    }
}

impl HTTPCompletionProvider for Google {
    fn complete_request(&self, req: &CompletionRequest) -> Result<Request<Vec<u8>>, LLMError> {
        let chat_message = ChatMessage::user().text(req.prompt.clone()).build();
        self.chat_request(&[chat_message], None)
    }

    fn parse_complete(&self, resp: Response<Vec<u8>>) -> Result<CompletionResponse, LLMError> {
        let chat_response = self.parse_chat(resp)?;
        if let Some(text) = chat_response.text() {
            Ok(CompletionResponse { text })
        } else {
            Err(LLMError::ProviderError(
                "No answer returned by Google".to_string(),
            ))
        }
    }
}

impl HTTPEmbeddingProvider for Google {
    fn embed_request(&self, inputs: &[String]) -> Result<Request<Vec<u8>>, LLMError> {
        let resolved_key = self.resolved_key();
        if resolved_key.is_empty() {
            return Err(LLMError::AuthError("Missing Google API key".to_string()));
        }
        let embedding_model = "text-embedding-004";

        //let mut embeddings = Vec::new();

        // Process each text separately as Gemini API accepts one text at a time
        let mut json_body;
        for text in inputs {
            let req_body = GoogleEmbeddingRequest {
                model: "models/text-embedding-004",
                content: GoogleEmbeddingContent {
                    parts: vec![GoogleContentPart::text(text)],
                },
            };
            json_body = serde_json::to_vec(&req_body)?;
        }

        let mut url = Google::default_base_url()
            .join(embedding_model)
            .map_err(|e| LLMError::HttpError(e.to_string()))?
            .join(":embedContent")
            .map_err(|e| LLMError::HttpError(e.to_string()))?;
        url.set_query(Some(&format!("key={}", &resolved_key)));

        unimplemented!();
        Err(LLMError::ProviderError("asd".to_string()))
        /*
        Ok(Request::builder()
            .method(Method::POST)
            .uri(url.as_str())
            .header(CONTENT_TYPE, "application/json")
            .body(json_body)?)
            */
    }

    fn parse_embed(&self, resp: Response<Vec<u8>>) -> Result<Vec<Vec<f32>>, LLMError> {
        let embedding_resp: GoogleEmbeddingResponse = serde_json::from_slice(resp.body())?;
        let _x = embedding_resp.embedding.values;
        //Ok(embedding_resp.embedding.values)
        todo!("finish google embedding");
        Err(LLMError::ProviderError("asd".to_string()))
    }
}

impl HTTPLLMProvider for Google {
    fn tools(&self) -> Option<&[Tool]> {
        self.tools.as_deref()
    }

    fn key_resolver(&self) -> Option<&Arc<dyn ApiKeyResolver>> {
        self.key_resolver.as_ref()
    }

    fn set_key_resolver(&mut self, resolver: Arc<dyn ApiKeyResolver>) {
        self.key_resolver = Some(resolver);
    }
}

#[derive(Default)]
struct GoogleStreamParser {
    buffer: String,
}

impl ChatStreamParser for GoogleStreamParser {
    fn parse_chunk(&mut self, chunk: &[u8]) -> Result<Vec<querymt::chat::StreamChunk>, LLMError> {
        let text =
            std::str::from_utf8(chunk).map_err(|e| LLMError::GenericError(format!("{:#}", e)))?;

        self.buffer.push_str(text);

        let (extracted_chunks, bytes_consumed) = extract_complete_json_objects(&self.buffer)?;

        if bytes_consumed > 0 {
            self.buffer.drain(..bytes_consumed);
        }

        for chunk in &extracted_chunks {
            if matches!(chunk, querymt::chat::StreamChunk::Done { .. }) {
                self.buffer.clear();
                break;
            }
        }

        Ok(extracted_chunks)
    }
}

/// Extract complete JSON objects from a buffer containing Google's streaming array format
/// Returns (extracted StreamChunks, number of bytes consumed from buffer)
fn extract_complete_json_objects(
    buffer: &str,
) -> Result<(Vec<querymt::chat::StreamChunk>, usize), LLMError> {
    let _result_chunks: Vec<querymt::chat::StreamChunk> = Vec::new();
    let mut bytes_consumed = 0;

    // Strip leading whitespace and array opening bracket
    let trimmed = buffer.trim_start();
    let working_text = if let Some(stripped) = trimmed.strip_prefix('[') {
        bytes_consumed += buffer.len() - trimmed.len() + 1; // whitespace + '['
        stripped
    } else {
        trimmed
    };

    // Strip leading comma and whitespace (between array elements)
    let working_text = working_text.trim_start();
    if let Some(stripped) = working_text.strip_prefix(',') {
        bytes_consumed += 1;
        let working_text = stripped.trim_start();
        bytes_consumed +=
            working_text.as_ptr() as usize - (buffer.as_ptr() as usize + bytes_consumed);

        // Now try to parse JSON objects from the working text
        return try_parse_json_objects(buffer, bytes_consumed, working_text);
    }

    try_parse_json_objects(buffer, bytes_consumed, working_text)
}

fn try_parse_json_objects(
    original_buffer: &str,
    initial_offset: usize,
    text: &str,
) -> Result<(Vec<querymt::chat::StreamChunk>, usize), LLMError> {
    use serde_json::Deserializer;

    let mut result_chunks = Vec::new();
    let mut total_consumed = initial_offset;

    // Try to parse JSON objects using StreamDeserializer
    let mut deserializer = Deserializer::from_str(text).into_iter::<GoogleChatResponse>();

    while let Some(result) = deserializer.next() {
        match result {
            Ok(response) => {
                // Track how many bytes we consumed
                let byte_offset = deserializer.byte_offset();
                total_consumed = initial_offset + byte_offset;

                // Extract StreamChunks from this response
                let chunks = extract_google_stream_chunks(response);
                result_chunks.extend(chunks);
            }
            Err(_e) => {
                // Parse error - likely incomplete JSON
                // Don't consume any more bytes - leave the rest in the buffer
                break;
            }
        }
    }

    // Check if there's a trailing ] (end of array)
    if total_consumed < original_buffer.len() {
        let remaining = &original_buffer[total_consumed..];
        let trimmed_remaining = remaining.trim_start();
        if trimmed_remaining.starts_with(']') {
            // Consume the closing bracket and any whitespace before it
            total_consumed += remaining.len() - trimmed_remaining.len() + 1;
        }
    }

    Ok((result_chunks, total_consumed))
}

/// Extract StreamChunks from a GoogleChatResponse
fn extract_google_stream_chunks(response: GoogleChatResponse) -> Vec<querymt::chat::StreamChunk> {
    let mut chunks = Vec::new();

    if let Some(candidate) = response.candidates.first() {
        // Extract text from parts
        for (index, part) in candidate.content.parts.iter().enumerate() {
            if let Some(text) = &part.text {
                if !text.is_empty() {
                    if part.thought {
                        chunks.push(querymt::chat::StreamChunk::Thinking(text.clone()));
                    } else {
                        chunks.push(querymt::chat::StreamChunk::Text(text.clone()));
                    }
                }
            }

            // Extract tool calls
            if let Some(function_call) = &part.function_call {
                let id = if let Some(sig) = &part.thought_signature {
                    format!("call_{}:{}", function_call.name, sig)
                } else {
                    format!("call_{}", function_call.name)
                };

                chunks.push(querymt::chat::StreamChunk::ToolUseStart {
                    index,
                    id: id.clone(),
                    name: function_call.name.clone(),
                });

                chunks.push(querymt::chat::StreamChunk::ToolUseComplete {
                    index,
                    tool_call: querymt::ToolCall {
                        id,
                        call_type: "function".to_string(),
                        function: querymt::FunctionCall {
                            name: function_call.name.clone(),
                            arguments: serde_json::to_string(&function_call.args)
                                .unwrap_or_default(),
                        },
                    },
                });
            }
        }

        // Handle content-level function calls (older format)
        if let Some(fc) = &candidate.content.function_call {
            let id = format!("call_{}", fc.name);

            chunks.push(querymt::chat::StreamChunk::ToolUseStart {
                index: 0,
                id: id.clone(),
                name: fc.name.clone(),
            });
            chunks.push(querymt::chat::StreamChunk::ToolUseComplete {
                index: 0,
                tool_call: querymt::ToolCall {
                    id,
                    call_type: "function".to_string(),
                    function: querymt::FunctionCall {
                        name: fc.name.clone(),
                        arguments: serde_json::to_string(&fc.args).unwrap_or_default(),
                    },
                },
            });
        }

        if let Some(fcs) = &candidate.content.function_calls {
            for (index, fc) in fcs.iter().enumerate() {
                let id = format!("call_{}", fc.name);

                chunks.push(querymt::chat::StreamChunk::ToolUseStart {
                    index,
                    id: id.clone(),
                    name: fc.name.clone(),
                });
                chunks.push(querymt::chat::StreamChunk::ToolUseComplete {
                    index,
                    tool_call: querymt::ToolCall {
                        id,
                        call_type: "function".to_string(),
                        function: querymt::FunctionCall {
                            name: fc.name.clone(),
                            arguments: serde_json::to_string(&fc.args).unwrap_or_default(),
                        },
                    },
                });
            }
        }

        // Emit usage BEFORE Done so consumers that break on Done still
        // capture Usage.  Google includes usage only in the final response
        // alongside finish_reason.
        if let Some(usage) = response.usage {
            chunks.push(querymt::chat::StreamChunk::Usage(usage));
        }

        // Check for finish reason (only in final chunk)
        if let Some(finish_reason) = &candidate.finish_reason {
            let has_tool_calls = candidate
                .content
                .function_calls
                .as_ref()
                .is_some_and(|fcs| !fcs.is_empty());
            let finish_reason = if has_tool_calls {
                FinishReason::ToolCalls
            } else {
                match finish_reason.as_str() {
                    "STOP" => FinishReason::Stop,
                    "MAX_TOKENS" => FinishReason::Length,
                    "SAFETY"
                    | "IMAGE_SAFETY"
                    | "IMAGE_PROHIBITED_CONTENT"
                    | "BLOCKLIST"
                    | "SPII"
                    | "PROHIBITED_CONTENT" => FinishReason::ContentFilter,
                    "OTHER" | "IMAGE_OTHER" => FinishReason::Other,
                    "MALFORMED_FUNCTION_CALL" => FinishReason::Error,
                    _ => FinishReason::Unknown,
                }
            };
            chunks.push(querymt::chat::StreamChunk::Done { finish_reason });
        }
    }

    chunks
}

struct GoogleFactory;
impl HTTPLLMProviderFactory for GoogleFactory {
    fn name(&self) -> &str {
        "google"
    }

    fn api_key_name(&self) -> Option<String> {
        Some("GEMINI_API_KEY".into())
    }

    fn list_models_request(&self, cfg: &str) -> Result<Request<Vec<u8>>, LLMError> {
        let cfg: Value = serde_json::from_str(cfg)?;
        let mut base_url = match cfg.get("base_url").and_then(Value::as_str) {
            Some(base_url_str) => Url::parse(base_url_str)?,
            None => Google::default_base_url(),
        };

        match cfg.get("api_key").and_then(Value::as_str) {
            Some(api_key) => {
                base_url.set_query(Some(&format!("key={}", api_key)));

                Ok(Request::builder()
                    .method(Method::GET)
                    .header(CONTENT_TYPE, "application/json")
                    .uri(base_url.as_str())
                    .body(Vec::new())?)
            }
            None => Err(LLMError::AuthError("Missing Google API key".into())),
        }
    }

    fn parse_list_models(&self, resp: Response<Vec<u8>>) -> Result<Vec<String>, LLMError> {
        let resp_json: Value = serde_json::from_slice(resp.body())?;
        let arr = resp_json
            .get("models")
            .and_then(Value::as_array)
            .ok_or_else(|| LLMError::InvalidRequest("`models` missing or not an array".into()))?;

        let names = arr
            .iter()
            .filter_map(|m| m.get("name"))
            .filter_map(Value::as_str)
            .filter_map(|v| v.strip_prefix("models/"))
            .map(String::from)
            .collect();
        Ok(names)
    }

    fn config_schema(&self) -> String {
        let schema = schema_for!(Google);
        serde_json::to_string(&schema).expect("Google JSON Schema should always serialize")
    }

    fn from_config(&self, cfg: &str) -> Result<Box<dyn HTTPLLMProvider>, LLMError> {
        let provider: Google = serde_json::from_str(cfg)?;
        Ok(Box::new(provider))
    }
}

/// Creates a Google HTTP factory for direct static registration.
pub fn create_http_factory() -> Arc<dyn HTTPLLMProviderFactory> {
    Arc::new(GoogleFactory)
}

#[cfg(feature = "native")]
#[unsafe(no_mangle)]
pub extern "C" fn plugin_http_factory() -> *mut dyn HTTPLLMProviderFactory {
    Box::into_raw(Box::new(GoogleFactory)) as *mut _
}

#[cfg(feature = "extism")]
mod extism_exports {
    use super::{Google, GoogleFactory};
    use querymt_extism_macros::impl_extism_http_plugin;

    impl_extism_http_plugin! {
        config = Google,
        factory = GoogleFactory,
        name   = "google",
    }
}
