use crate::core::LlmClientTrait;
use crate::types::WorkerError;
use worker_types::{DynamicToolDefinition, LlmProvider, Message, Role, StreamEvent, ToolCall};
use crate::config::UrlConfig;
use futures_util::{Stream, StreamExt, TryStreamExt};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tracing;

/// Extract tool name from Tool message content
fn extract_tool_name_from_content(content: &str) -> Option<String> {
    // Look for patterns like "Tool 'tool_name' executed successfully"
    if let Some(start) = content.find("Tool '") {
        if let Some(end) = content[start + 6..].find("'") {
            let tool_name = &content[start + 6..start + 6 + end];
            return Some(tool_name.to_string());
        }
    }
    None
}

/// Parse tool call information from message content

/// Transforms a JSON schema to be compatible with Gemini API
/// Converts 'uint' types to 'integer' types and handles nullable types
/// Also ensures the schema is in the correct format for Gemini function parameters
fn transform_schema_for_gemini(schema: serde_json::Value) -> serde_json::Value {
    match schema {
        serde_json::Value::Object(mut obj) => {
            // Remove $schema key as it's not needed for Gemini
            obj.remove("$schema");

            // Handle type field
            if let Some(type_val) = obj.get("type") {
                match type_val {
                    // Convert 'uint' to 'integer'
                    serde_json::Value::String(s) if s == "uint" => {
                        obj.insert(
                            "type".to_string(),
                            serde_json::Value::String("integer".to_string()),
                        );
                        // Add format for integer types as required by Gemini
                        obj.insert(
                            "format".to_string(),
                            serde_json::Value::String("int64".to_string()),
                        );
                    }
                    // Handle array types like ["integer", "null"]
                    serde_json::Value::Array(arr) => {
                        if let Some(non_null_type) = arr.iter().find(|&t| t != "null") {
                            // Use the non-null type
                            let mut new_type = non_null_type.clone();
                            // Convert 'uint' to 'integer' if needed
                            if let serde_json::Value::String(s) = &new_type {
                                if s == "uint" {
                                    new_type = serde_json::Value::String("integer".to_string());
                                }
                            }
                            obj.insert("type".to_string(), new_type.clone());

                            // Add format for integer types as required by Gemini
                            if let serde_json::Value::String(type_str) = &new_type {
                                if type_str == "integer" {
                                    obj.insert(
                                        "format".to_string(),
                                        serde_json::Value::String("int64".to_string()),
                                    );
                                }
                            }
                        }
                    }
                    // Handle existing integer types
                    serde_json::Value::String(s) if s == "integer" => {
                        // Add format for integer types as required by Gemini
                        obj.insert(
                            "format".to_string(),
                            serde_json::Value::String("int64".to_string()),
                        );
                    }
                    _ => {}
                }
            }

            // Handle properties and required fields
            if let (Some(properties), Some(required)) = (obj.get("properties"), obj.get("required"))
            {
                if let (serde_json::Value::Object(props), serde_json::Value::Array(req_arr)) =
                    (properties, required)
                {
                    let mut new_required = Vec::new();

                    for (prop_name, _) in props {
                        // Only include in required if it's not nullable
                        if req_arr.iter().any(|r| r == prop_name) {
                            // Check if this property has a nullable type
                            if let Some(prop_schema) = props.get(prop_name) {
                                if let Some(prop_type) = prop_schema.get("type") {
                                    // If type is an array containing "null", it's nullable
                                    let is_nullable = match prop_type {
                                        serde_json::Value::Array(arr) => {
                                            arr.iter().any(|t| t == "null")
                                        }
                                        _ => false,
                                    };

                                    // Only add to required if not nullable
                                    if !is_nullable {
                                        new_required
                                            .push(serde_json::Value::String(prop_name.clone()));
                                    }
                                } else {
                                    // No type info, assume required
                                    new_required.push(serde_json::Value::String(prop_name.clone()));
                                }
                            }
                        }
                    }

                    obj.insert(
                        "required".to_string(),
                        serde_json::Value::Array(new_required),
                    );
                }
            }

            // Recursively transform nested objects
            for (_, value) in obj.iter_mut() {
                *value = transform_schema_for_gemini(value.clone());
            }

            serde_json::Value::Object(obj)
        }
        serde_json::Value::Array(arr) => {
            serde_json::Value::Array(arr.into_iter().map(transform_schema_for_gemini).collect())
        }
        other => other,
    }
}

// --- Request Structures ---
#[derive(Debug, Serialize, Clone)]
pub struct GeminiTool {
    #[serde(rename = "functionDeclarations")]
    pub function_declarations: Vec<GeminiFunctionDeclaration>,
}

#[derive(Debug, Serialize, Clone)]
pub struct GeminiFunctionDeclaration {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

#[derive(Debug, Serialize, Clone)]
pub struct GeminiRequest {
    pub contents: Vec<GeminiContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(rename = "systemInstruction")]
    pub system_instruction: Option<GeminiContent>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<GeminiTool>>,
}

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

#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub enum GeminiPart {
    Text {
        text: String,
    },
    FunctionCall {
        #[serde(rename = "functionCall")]
        function_call: GeminiFunctionCall,
    },
    FunctionResponse {
        #[serde(rename = "functionResponse")]
        function_response: GeminiFunctionResponse,
    },
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct GeminiFunctionCall {
    pub name: String,
    pub args: serde_json::Value,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct GeminiFunctionResponse {
    pub name: String,
    pub response: serde_json::Value,
}

// --- Response Structures ---
#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct GeminiResponse {
    #[serde(default)]
    pub candidates: Vec<GeminiCandidate>,
}

#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct GeminiCandidate {
    pub content: GeminiContent,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finish_reason: Option<String>,
}

fn build_url(model: &str) -> String {
    let base_url = UrlConfig::get_base_url("gemini");
    let action = "streamGenerateContent";
    format!("{}/v1beta/models/{}:{}", base_url, model, action)
}

/// Finds the start and end indices of the first complete JSON object `{...}` in the buffer.
fn find_first_json_object_bounds(buffer: &[u8]) -> Option<(usize, usize)> {
    if let Some(start) = buffer.iter().position(|&b| b == b'{') {
        let mut brace_count = 0;
        let mut in_string = false;
        let mut escaped = false;

        for (i, &byte) in buffer.iter().skip(start).enumerate() {
            if in_string {
                if escaped {
                    escaped = false;
                } else if byte == b'\\' {
                    escaped = true;
                } else if byte == b'"' {
                    in_string = false;
                }
            } else {
                match byte {
                    b'"' => in_string = true,
                    b'{' => brace_count += 1,
                    b'}' => {
                        brace_count -= 1;
                        if brace_count == 0 {
                            let end = start + i + 1;
                            return Some((start, end));
                        }
                    }
                    _ => {}
                }
            }
        }
    }
    None // No complete object found
}

/// Completes a chat request with streaming, yielding StreamEvent objects.
pub(crate) fn stream_events<'a>(
    api_key: &'a str,
    model: &'a str,
    request: GeminiRequest,
    llm_debug: Option<crate::types::LlmDebug>,
) -> impl Stream<Item = anyhow::Result<StreamEvent>> + 'a {
    let api_key = api_key.to_string();
    let model = model.to_string();

    async_stream::try_stream! {
        let body = serde_json::to_string_pretty(&request).unwrap_or_else(|e| e.to_string());
        tracing::debug!("Gemini Request Body: {}", body);

        if let Some(debug_settings) = &llm_debug {
            if let Some(debug_event) = debug_settings.debug_request(&model, "Gemini", &serde_json::to_value(&request).unwrap_or_default()) {
                yield debug_event;
            }
        }

        let client = Client::new();
        let url = build_url(&model);

        let response = client
            .post(&url)
            .header("x-goog-api-key", &api_key)
            .json(&request)
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("Gemini API request failed: {}", e))?;

        let status = response.status();
        if !status.is_success() {
            let error_body = response.text().await.unwrap_or_else(|_| "Could not read error body".to_string());
            let error_msg = format!("Gemini API request failed with status: {} - {}", status, error_body);
            tracing::error!("{}", error_msg);
            Err(anyhow::anyhow!(error_msg))?;
        } else {
            let mut byte_stream = response.bytes_stream();
            let mut buffer = Vec::new();
            let mut full_content = String::new();

            while let Some(chunk_result) = byte_stream.next().await {
                let chunk = chunk_result?;
                buffer.extend_from_slice(&chunk);

                while let Some((start, end)) = find_first_json_object_bounds(&buffer) {
                    let json_slice = &buffer[start..end];

                    if let Some(debug_settings) = &llm_debug {
                        if let Ok(response_value) = serde_json::from_slice::<serde_json::Value>(json_slice) {
                            if let Some(debug_event) = debug_settings.debug_response(&model, "Gemini", &response_value) {
                                yield debug_event;
                            }
                        }
                    }

                    match serde_json::from_slice::<GeminiResponse>(json_slice) {
                        Ok(response) => {
                            let response_text = String::from_utf8_lossy(json_slice);
                            tracing::debug!(
                                response = %response_text,
                                candidates_count = response.candidates.len(),
                                "Successfully parsed Gemini response"
                            );
                            if response.candidates.is_empty() {
                                tracing::warn!(
                                    response = %response_text,
                                    "Received empty candidates in Gemini response"
                                );
                            } else if let Some(candidate) = response.candidates.get(0) {
                                // Log finish reason for debugging
                                if let Some(ref finish_reason) = candidate.finish_reason {
                                    tracing::debug!(
                                        finish_reason = %finish_reason,
                                        "Received finish reason in Gemini response"
                                    );

                                    // Handle specific finish reasons
                                    match finish_reason.as_str() {
                                        "STOP" => {
                                            tracing::debug!("Gemini response completed with STOP");
                                            // Continue processing parts if any, this is normal completion
                                        }
                                        "MAX_TOKENS" => {
                                            tracing::warn!("Gemini response stopped due to MAX_TOKENS");
                                        }
                                        "SAFETY" => {
                                            tracing::warn!("Gemini response stopped due to SAFETY concerns");
                                        }
                                        "RECITATION" => {
                                            tracing::warn!("Gemini response stopped due to RECITATION");
                                        }
                                        other => {
                                            tracing::warn!("Gemini response stopped with unknown reason: {}", other);
                                        }
                                    }
                                }

                                if candidate.content.parts.is_empty() {
                                    tracing::warn!(
                                        response = %response_text,
                                        role = %candidate.content.role,
                                        finish_reason = ?candidate.finish_reason,
                                        "Received empty parts in Gemini response"
                                    );
                                } else {
                                    for part in &candidate.content.parts {
                                        tracing::debug!("Processing Gemini part (type unknown)");
                                        match part {
                                            GeminiPart::Text { text } => {
                                                tracing::debug!("Found Text part with content length: {}", text.len());
                                                full_content.push_str(text);
                                                yield StreamEvent::Chunk(text.clone());
                                            }
                                            GeminiPart::FunctionCall { function_call } => {
                                                tracing::debug!("Found FunctionCall part: name={}, args={:?}", function_call.name, function_call.args);
                                                let tool_call = ToolCall {
                                                    name: function_call.name.clone(),
                                                    arguments: serde_json::to_string(&function_call.args)
                                                        .unwrap_or_else(|_| "{}".to_string()),
                                                };
                                                yield StreamEvent::ToolCall(tool_call.clone());
                                            }
                                            GeminiPart::FunctionResponse { .. } => {
                                                // Function responses in model output are not expected
                                                // as they're part of the input conversation history
                                                tracing::warn!("Unexpected FunctionResponse in model output");
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        Err(e) => {
                            let response_text = String::from_utf8_lossy(json_slice);
                            tracing::warn!(
                                error = %e,
                                response = %response_text,
                                "Failed to deserialize GeminiResponse from slice"
                            );
                        }
                    }
                    buffer.drain(..end);
                }
            }

            let final_message = Message::new(
                Role::Model,
                full_content,
            );
            yield StreamEvent::Completion(final_message);
        }
    }
}

pub struct GeminiClient {
    api_key: String,
    model: String,
}

impl GeminiClient {
    pub fn new(api_key: &str, model: &str) -> Self {
        Self {
            api_key: api_key.to_string(),
            model: model.to_string(),
        }
    }

    pub fn get_model_name(&self) -> String {
        self.model.clone()
    }

    /// 静的メソッド：API キーを受け取ってモデル一覧を取得
    pub async fn list_models_static(
        api_key: &str,
    ) -> Result<Vec<crate::types::ModelInfo>, WorkerError> {
        let client = Client::new();
        let url = UrlConfig::get_models_url("gemini");

        let response = client
            .get(url)
            .header("x-goog-api-key", api_key)
            .send()
            .await
            .map_err(|e| {
                tracing::error!("Gemini API request failed: {}", e);
                WorkerError::from_api_error(e.to_string(), &crate::types::LlmProvider::Gemini)
            })?;

        if !response.status().is_success() {
            let status = response.status();
            let error_body = response.text().await.unwrap_or_default();
            tracing::error!(
                "Gemini list_models_static failed - Status: {}, Body: {}",
                status,
                error_body
            );
            return Err(WorkerError::from_api_error(
                format!("Failed to list Gemini models: {} - {}", status, error_body),
                &crate::types::LlmProvider::Gemini,
            ));
        }

        let models_response: serde_json::Value = response.json().await.map_err(|e| {
            WorkerError::from_api_error(e.to_string(), &crate::types::LlmProvider::Gemini)
        })?;

        let mut models = Vec::new();

        if let Some(models_array) = models_response.get("models").and_then(|m| m.as_array()) {
            for model in models_array {
                if let Some(name) = model.get("name").and_then(|n| n.as_str()) {
                    // "models/" プレフィックスを除去
                    let model_id = name.strip_prefix("models/").unwrap_or(name);

                    // generateContentメソッドをサポートするモデルのみを含める
                    if let Some(supported_methods) = model
                        .get("supportedGenerationMethods")
                        .and_then(|m| m.as_array())
                    {
                        let supports_generate_content = supported_methods
                            .iter()
                            .any(|method| method.as_str() == Some("generateContent"));

                        if supports_generate_content {
                            models.push(crate::types::ModelInfo {
                                id: model_id.to_string(),
                                name: model
                                    .get("displayName")
                                    .and_then(|d| d.as_str())
                                    .unwrap_or(model_id)
                                    .to_string(),
                                provider: crate::types::LlmProvider::Gemini,
                                supports_tools: true,
                                supports_function_calling: true,
                                supports_vision: false,
                                supports_multimodal: false,
                                context_length: None,
                                training_cutoff: None,
                                capabilities: vec!["text_generation".to_string()],
                                description: model
                                    .get("description")
                                    .and_then(|d| d.as_str())
                                    .map(|s| s.to_string())
                                    .or_else(|| Some(format!("Google Gemini model: {}", model_id))),
                            });
                        }
                    }
                }
            }
        }

        tracing::info!(
            "Gemini list_models_static found {} models with metadata",
            models.len()
        );
        Ok(models)
    }

    pub async fn chat_stream<'a>(
        &'a self,
        messages: Vec<Message>,
        tools: Option<&[DynamicToolDefinition]>,
        llm_debug: Option<crate::types::LlmDebug>,
    ) -> Result<
        Box<dyn Stream<Item = Result<StreamEvent, WorkerError>> + Unpin + Send + 'a>,
        WorkerError,
    > {
        // Separate system messages from regular messages
        let (system_messages, regular_messages): (Vec<_>, Vec<_>) = messages
            .into_iter()
            .partition(|msg| matches!(msg.role, Role::System));

        // Create system instruction from system messages
        let system_instruction = if !system_messages.is_empty() {
            let combined_system_content = system_messages
                .into_iter()
                .map(|msg| msg.content)
                .collect::<Vec<_>>()
                .join("\n\n");

            Some(GeminiContent {
                role: "user".to_string(), // System instruction uses "user" role
                parts: vec![GeminiPart::Text {
                    text: combined_system_content,
                }],
            })
        } else {
            None
        };

        // Process regular messages with proper tool context handling
        let contents = regular_messages
            .into_iter()
            .map(|msg| {
                let (role, parts) = match msg.role {
                    Role::User => (
                        "user".to_string(),
                        vec![GeminiPart::Text { text: msg.content }],
                    ),
                    Role::Model => {
                        if let Some(tool_calls) = &msg.tool_calls {
                            // Model message with tool calls - convert to FunctionCall parts
                            tracing::debug!(
                                "Converting model message with {} tool calls to FunctionCall parts",
                                tool_calls.len()
                            );
                            let mut parts = Vec::new();

                            // Add text content if present
                            if !msg.content.is_empty() {
                                parts.push(GeminiPart::Text {
                                    text: msg.content.clone(),
                                });
                            }

                            // Add function calls
                            for tool_call in tool_calls {
                                tracing::debug!(
                                    "Adding FunctionCall part for tool: {}",
                                    tool_call.name
                                );
                                let args = serde_json::from_str(&tool_call.arguments)
                                    .unwrap_or(serde_json::json!({}));
                                parts.push(GeminiPart::FunctionCall {
                                    function_call: GeminiFunctionCall {
                                        name: tool_call.name.clone(),
                                        args,
                                    },
                                });
                            }

                            ("model".to_string(), parts)
                        } else {
                            // Regular model message
                            tracing::debug!("Converting regular model message (no tool calls)");
                            (
                                "model".to_string(),
                                vec![GeminiPart::Text { text: msg.content }],
                            )
                        }
                    }
                    Role::Tool => {
                        // Tool responses should be sent as FunctionResponse
                        if let Some(tool_name) = extract_tool_name_from_content(&msg.content) {
                            // Extract result from the content
                            let result_value = if msg.content.contains("Result: ") {
                                if let Some(result_start) = msg.content.find("Result: ") {
                                    let result_str = &msg.content[result_start + 8..];
                                    // Try to parse as JSON, fallback to string
                                    serde_json::from_str(result_str)
                                        .unwrap_or_else(|_| serde_json::json!(result_str))
                                } else {
                                    serde_json::json!(msg.content)
                                }
                            } else {
                                serde_json::json!(msg.content)
                            };

                            (
                                "user".to_string(),
                                vec![GeminiPart::FunctionResponse {
                                    function_response: GeminiFunctionResponse {
                                        name: tool_name,
                                        response: result_value,
                                    },
                                }],
                            )
                        } else {
                            // Fallback to text response if tool name can't be extracted
                            (
                                "user".to_string(),
                                vec![GeminiPart::Text {
                                    text: format!("Tool Response:\n{}", msg.content),
                                }],
                            )
                        }
                    }
                    Role::System => unreachable!(), // Should not reach here after partition
                };
                GeminiContent { role, parts }
            })
            .collect();

        let tools = tools.map(|tools| {
            vec![GeminiTool {
                function_declarations: tools
                    .iter()
                    .map(|tool| {
                        let mut transformed_schema =
                            transform_schema_for_gemini(tool.parameters_schema.clone());

                        // Ensure the schema has the correct structure for Gemini
                        match transformed_schema {
                            serde_json::Value::Object(ref mut obj) => {
                                // Gemini expects the parameters to be an object with type: "object"
                                if !obj.contains_key("type") {
                                    obj.insert(
                                        "type".to_string(),
                                        serde_json::Value::String("object".to_string()),
                                    );
                                }
                                // If there are no properties, add an empty object
                                if !obj.contains_key("properties") {
                                    obj.insert(
                                        "properties".to_string(),
                                        serde_json::Value::Object(serde_json::Map::new()),
                                    );
                                }
                                // If there are no required fields, add an empty array
                                if !obj.contains_key("required") {
                                    obj.insert(
                                        "required".to_string(),
                                        serde_json::Value::Array(vec![]),
                                    );
                                }
                            }
                            _ => {
                                // If it's not an object, create a proper object schema
                                let mut schema_obj = serde_json::Map::new();
                                schema_obj.insert(
                                    "type".to_string(),
                                    serde_json::Value::String("object".to_string()),
                                );
                                schema_obj.insert(
                                    "properties".to_string(),
                                    serde_json::Value::Object(serde_json::Map::new()),
                                );
                                schema_obj.insert(
                                    "required".to_string(),
                                    serde_json::Value::Array(vec![]),
                                );
                                transformed_schema = serde_json::Value::Object(schema_obj);
                            }
                        }

                        GeminiFunctionDeclaration {
                            name: tool.name.clone(),
                            description: tool.description.clone(),
                            parameters: transformed_schema,
                        }
                    })
                    .collect(),
            }]
        });

        let request = GeminiRequest {
            contents,
            system_instruction,
            tools,
        };

        let stream = stream_events(&self.api_key, &self.model, request, llm_debug)
            .map_err(|e| WorkerError::llm_api("gemini", e.to_string()));

        Ok(Box::new(Box::pin(stream)))
    }

    pub async fn get_model_details(
        &self,
        model_name: &str,
    ) -> Result<crate::types::ModelInfo, WorkerError> {
        let client = Client::new();
        let url = UrlConfig::get_model_url("gemini", model_name);

        let response = client
            .get(&url)
            .header("x-goog-api-key", &self.api_key)
            .send()
            .await
            .map_err(|e| {
                WorkerError::from_api_error(e.to_string(), &crate::types::LlmProvider::Gemini)
            })?;

        if !response.status().is_success() {
            return Err(WorkerError::from_api_error(
                format!(
                    "Gemini model details request failed with status: {}",
                    response.status()
                ),
                &crate::types::LlmProvider::Gemini,
            ));
        }

        let model_data: serde_json::Value = response.json().await.map_err(|e| {
            WorkerError::from_api_error(e.to_string(), &crate::types::LlmProvider::Gemini)
        })?;

        let name = model_data
            .get("name")
            .and_then(|n| n.as_str())
            .unwrap_or(model_name);
        let display_name = model_data
            .get("displayName")
            .and_then(|d| d.as_str())
            .unwrap_or(name);
        let description = model_data
            .get("description")
            .and_then(|d| d.as_str())
            .unwrap_or("");
        let version = model_data
            .get("version")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let input_token_limit = model_data
            .get("inputTokenLimit")
            .and_then(|i| i.as_u64())
            .map(|i| i as u32);
        let _output_token_limit = model_data
            .get("outputTokenLimit")
            .and_then(|o| o.as_u64())
            .map(|o| o as u32);

        let empty_vec = Vec::new();
        let supported_methods = model_data
            .get("supportedGenerationMethods")
            .and_then(|s| s.as_array())
            .unwrap_or(&empty_vec);

        let supports_tools = supported_methods
            .iter()
            .any(|method| method.as_str() == Some("generateContent"));

        let supports_vision = false; // Will be determined dynamically
        let capabilities = vec!["text_generation".to_string()]; // Basic default

        Ok(crate::types::ModelInfo {
            id: model_name.to_string(),
            name: display_name.to_string(),
            provider: crate::types::LlmProvider::Gemini,
            supports_tools,
            supports_function_calling: supports_tools,
            supports_vision,
            supports_multimodal: supports_vision,
            context_length: input_token_limit,
            training_cutoff: version,
            capabilities,
            description: Some(if description.is_empty() {
                format!("Google Gemini model: {}", display_name)
            } else {
                description.to_string()
            }),
        })
    }

    pub async fn check_connection(&self) -> Result<(), WorkerError> {
        // Simple connection check - try to call the API
        // For now, just return OK if model is not empty
        if self.model.is_empty() {
            return Err(WorkerError::model_not_found("gemini", "No model specified"));
        }
        Ok(())
    }
}

#[async_trait::async_trait]
impl LlmClientTrait for GeminiClient {
    async fn chat_stream<'a>(
        &'a self,
        messages: Vec<Message>,
        tools: Option<&[DynamicToolDefinition]>,
        llm_debug: Option<crate::types::LlmDebug>,
    ) -> Result<
        Box<dyn Stream<Item = Result<StreamEvent, WorkerError>> + Unpin + Send + 'a>,
        WorkerError,
    > {
        self.chat_stream(messages, tools, llm_debug).await
    }

    async fn check_connection(&self) -> Result<(), WorkerError> {
        self.check_connection().await
    }

    fn provider(&self) -> LlmProvider {
        LlmProvider::Gemini
    }

    fn get_model_name(&self) -> String {
        self.get_model_name()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::env;

    #[test]
    fn test_schema_transformation() {
        // Test schema with various type formats including $schema
        let schema = serde_json::json!({
            "$schema": "https://json-schema.org/draft/2020-12/schema",
            "type": "object",
            "properties": {
                "id": {
                    "type": "uint"
                },
                "optional_number": {
                    "type": ["integer", "null"]
                },
                "required_string": {
                    "type": "string"
                },
                "existing_integer": {
                    "type": "integer"
                },
                "nested": {
                    "type": "object",
                    "properties": {
                        "count": {
                            "type": ["uint", "null"]
                        }
                    }
                }
            },
            "required": ["id", "optional_number", "required_string"]
        });

        let transformed = transform_schema_for_gemini(schema);

        // Check that the schema has the correct structure
        assert_eq!(transformed["type"], "object");
        assert!(transformed["properties"].is_object());
        assert!(transformed["required"].is_array());

        // Check that $schema key is removed
        assert!(transformed.get("$schema").is_none());

        // Check that 'uint' was transformed to 'integer'
        assert_eq!(transformed["properties"]["id"]["type"], "integer");
        assert_eq!(transformed["properties"]["id"]["format"], "int64");

        // Check that array types are converted to single types
        assert_eq!(
            transformed["properties"]["optional_number"]["type"],
            "integer"
        );
        assert_eq!(
            transformed["properties"]["optional_number"]["format"],
            "int64"
        );
        assert_eq!(
            transformed["properties"]["nested"]["properties"]["count"]["type"],
            "integer"
        );
        assert_eq!(
            transformed["properties"]["nested"]["properties"]["count"]["format"],
            "int64"
        );

        // Check that existing integer types also get format
        assert_eq!(
            transformed["properties"]["existing_integer"]["type"],
            "integer"
        );
        assert_eq!(
            transformed["properties"]["existing_integer"]["format"],
            "int64"
        );

        // Check that required array is updated correctly (nullable properties should be removed)
        let required: Vec<&str> = transformed["required"]
            .as_array()
            .unwrap()
            .iter()
            .map(|v| v.as_str().unwrap())
            .collect();

        assert!(required.contains(&"id"));
        assert!(required.contains(&"required_string"));
        assert!(!required.contains(&"optional_number")); // Should be removed because it's nullable
    }

    #[test]
    fn test_empty_schema_transformation() {
        // Test with an empty schema as would be processed in tool generation
        let schema = serde_json::json!({});
        let mut transformed = transform_schema_for_gemini(schema);

        // Apply the same logic as in tool generation
        match transformed {
            serde_json::Value::Object(ref mut obj) => {
                if !obj.contains_key("type") {
                    obj.insert(
                        "type".to_string(),
                        serde_json::Value::String("object".to_string()),
                    );
                }
                if !obj.contains_key("properties") {
                    obj.insert(
                        "properties".to_string(),
                        serde_json::Value::Object(serde_json::Map::new()),
                    );
                }
                if !obj.contains_key("required") {
                    obj.insert("required".to_string(), serde_json::Value::Array(vec![]));
                }
            }
            _ => {
                let mut schema_obj = serde_json::Map::new();
                schema_obj.insert(
                    "type".to_string(),
                    serde_json::Value::String("object".to_string()),
                );
                schema_obj.insert(
                    "properties".to_string(),
                    serde_json::Value::Object(serde_json::Map::new()),
                );
                schema_obj.insert("required".to_string(), serde_json::Value::Array(vec![]));
                transformed = serde_json::Value::Object(schema_obj);
            }
        }

        // Should be converted to a proper object schema
        assert_eq!(transformed["type"], "object");
        assert!(transformed["properties"].is_object());
        assert!(transformed["required"].is_array());
    }
}