//! Google (Gemini) provider impl.
//!
//! Wraps the Gemini REST API at
//! `https://generativelanguage.googleapis.com/v1beta`. Two endpoints
//! matter: `models/{model}:generateContent` for completions and
//! `models/{model}:countTokens` for token counts. Auth is the
//! `x-goog-api-key` header so the same client works against either
//! AI Studio or a Workspace-issued key. JSON mode rides on
//! `generationConfig.responseMimeType = "application/json"`; an
//! optional `responseSchema` constrains the shape further.

use std::time::Duration;

use async_trait::async_trait;
use serde::Deserialize;
use serde_json::{json, Value};

use crate::eval::error::ProviderError;
use crate::eval::provider::{
    map_reqwest_error, Capabilities, CompletionOptions, Cost, FinishReason, LlmProvider, Message,
    ProviderId, Response, ResponseContent, RetryPolicy, Role, ToolCall, ToolChoice, Usage,
};
use crate::eval::providers::parse_embedding;
use crate::eval::providers::retry::run_with_retry;

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

/// Google Gemini provider client.
///
/// Construct via [`GoogleProvider::new`]. Holds a `reqwest::Client` that
/// reuses connections across calls; provider impls are designed to be
/// long-lived (one per juror).
pub struct GoogleProvider {
    client: reqwest::Client,
    api_key: String,
    base_url: String,
    retry: RetryPolicy,
    request_timeout: Duration,
}

/// Config bag passed to [`GoogleProvider::new`].
///
/// `api_key` is resolved by the loader from `${env:GOOGLE_API_KEY}`; the
/// loader fails fast when the env var is missing so a CI run never
/// reaches the network with an empty header.
#[derive(Debug, Clone)]
pub struct GoogleConfig {
    /// API key, already resolved from the env.
    pub api_key: String,
    /// Override the base URL (Vertex gateways and `wiremock` tests).
    pub base_url: Option<String>,
    /// HTTP request timeout. Default 30s.
    pub request_timeout: Duration,
    /// HTTP / HTTPS proxy policy.
    pub proxy: crate::network::ProxyConfig,
}

impl Default for GoogleConfig {
    fn default() -> Self {
        Self {
            api_key: String::new(),
            base_url: None,
            request_timeout: Duration::from_secs(30),
            proxy: crate::network::ProxyConfig::default(),
        }
    }
}

impl GoogleProvider {
    /// Build a new provider client. Returns [`ProviderError::Transport`]
    /// when the underlying `reqwest::Client` cannot be constructed (TLS
    /// initialization failure).
    pub fn new(config: GoogleConfig) -> Result<Self, ProviderError> {
        let builder = reqwest::Client::builder().timeout(config.request_timeout);
        let builder = config
            .proxy
            .apply(builder)
            .map_err(|e| ProviderError::Transport(e.to_string()))?;
        let client = builder
            .build()
            .map_err(|e| ProviderError::Transport(e.to_string()))?;
        Ok(Self {
            client,
            api_key: config.api_key,
            base_url: config.base_url.unwrap_or_else(|| DEFAULT_BASE_URL.into()),
            retry: RetryPolicy::default(),
            request_timeout: config.request_timeout,
        })
    }
}

#[async_trait]
impl LlmProvider for GoogleProvider {
    fn name(&self) -> ProviderId {
        ProviderId::Google
    }

    fn supported_capabilities(&self) -> Capabilities {
        Capabilities {
            json_mode: true,
            tool_use: true,
            streaming: true,
            structured_output: true,
            // `countTokens` returns an exact count from the same
            // tokenizer the model uses for billing.
            token_count_exact: true,
        }
    }

    async fn chat_completion(
        &self,
        messages: &[Message],
        options: &CompletionOptions,
    ) -> Result<Response, ProviderError> {
        let body = build_generate_body(messages, options);
        let url = format!("{}/models/{}:generateContent", self.base_url, options.model);
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .header("x-goog-api-key", &self.api_key)
                .json(&body)
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        parse_generate_response(response_json, &options.model)
    }

    async fn count_tokens(&self, text: &str, model: &str) -> Result<u32, ProviderError> {
        let url = format!("{}/models/{}:countTokens", self.base_url, model);
        let body = json!({
            "contents": [{
                "role": "user",
                "parts": [{"text": text}],
            }],
        });
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .header("x-goog-api-key", &self.api_key)
                .json(&body)
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        response_json
            .get("totalTokens")
            .and_then(|v| v.as_u64())
            .map(|n| n as u32)
            .ok_or_else(|| {
                ProviderError::MalformedResponse("countTokens missing totalTokens".into())
            })
    }

    async fn embed(&self, text: &str, model: &str) -> Result<Vec<f32>, ProviderError> {
        let url = format!("{}/models/{}:embedContent", self.base_url, model);
        let body = json!({
            "model": format!("models/{model}"),
            "content": { "parts": [{ "text": text }] },
        });
        let response_json: Value = run_with_retry(&self.retry, || async {
            let resp = self
                .client
                .post(&url)
                .header("x-goog-api-key", &self.api_key)
                .json(&body)
                .send()
                .await
                .map_err(|e| map_reqwest_error(e, self.request_timeout))?;
            handle_response(resp).await
        })
        .await?;
        parse_embedding(&response_json, "/embedding/values")
    }

    fn estimate_cost(&self, usage: &Usage, model: &str) -> Cost {
        let rates = google_rates(model);
        Cost {
            input_dollars: rates.input_per_token * f64::from(usage.input_tokens),
            output_dollars: rates.output_per_token * f64::from(usage.output_tokens),
            cache_read_dollars: None,
            cache_write_dollars: None,
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct GoogleRates {
    input_per_token: f64,
    output_per_token: f64,
}

/// Published pricing (USD per token) for the four currently-supported
/// Gemini models. Unknown model strings fall back to the cheapest tier
/// (`gemini-1.5-flash`) so a typo does not surface a million-dollar
/// estimate.
fn google_rates(model: &str) -> GoogleRates {
    let lower = model.to_ascii_lowercase();
    if lower.starts_with("gemini-2.0-pro") {
        // Gemini 2.0 Pro: $1.25 / $5 per million tokens.
        GoogleRates {
            input_per_token: 1.25 / 1_000_000.0,
            output_per_token: 5.00 / 1_000_000.0,
        }
    } else if lower.starts_with("gemini-2.0-flash") {
        // Gemini 2.0 Flash: $0.10 / $0.40 per million tokens.
        GoogleRates {
            input_per_token: 0.10 / 1_000_000.0,
            output_per_token: 0.40 / 1_000_000.0,
        }
    } else if lower.starts_with("gemini-1.5-pro") {
        // Gemini 1.5 Pro: $1.25 / $5 per million tokens (prompts under 128k).
        GoogleRates {
            input_per_token: 1.25 / 1_000_000.0,
            output_per_token: 5.00 / 1_000_000.0,
        }
    } else {
        // gemini-1.5-flash (default) at $0.075 / $0.30 per million.
        GoogleRates {
            input_per_token: 0.075 / 1_000_000.0,
            output_per_token: 0.30 / 1_000_000.0,
        }
    }
}

fn build_generate_body(messages: &[Message], options: &CompletionOptions) -> Value {
    let (system, chat): (Vec<&Message>, Vec<&Message>) = messages
        .iter()
        .partition(|m| matches!(m.role, Role::System));
    let system_text = system
        .iter()
        .map(|m| m.content.as_str())
        .collect::<Vec<_>>()
        .join("\n\n");
    let contents: Vec<Value> = chat
        .iter()
        .map(|m| {
            let role = match m.role {
                Role::Assistant => "model",
                // Gemini accepts only `user` and `model` roles; tool
                // results ride alongside the user turn that followed.
                _ => "user",
            };
            json!({
                "role": role,
                "parts": [{"text": m.content}],
            })
        })
        .collect();
    let mut body = json!({
        "contents": contents,
    });
    if !system_text.is_empty() {
        body["systemInstruction"] = json!({
            "parts": [{"text": system_text}],
        });
    }
    let mut generation_config = json!({});
    if let Some(t) = options.temperature {
        generation_config["temperature"] = json!(t);
    }
    if let Some(n) = options.max_tokens {
        generation_config["maxOutputTokens"] = json!(n);
    }
    if options.json_mode {
        generation_config["responseMimeType"] = json!("application/json");
    }
    if generation_config
        .as_object()
        .map(|o| !o.is_empty())
        .unwrap_or(false)
    {
        body["generationConfig"] = generation_config;
    }
    if !options.tools.is_empty() {
        let declarations: Vec<Value> = options
            .tools
            .iter()
            .map(|t| {
                json!({
                    "name": t.name,
                    "description": t.description,
                    "parameters": gemini_schema(&t.input_schema),
                })
            })
            .collect();
        body["tools"] = json!([{"functionDeclarations": declarations}]);
        body["toolConfig"] = match &options.tool_choice {
            Some(ToolChoice::Tool(name)) => json!({
                "functionCallingConfig": {
                    "mode": "ANY",
                    "allowedFunctionNames": [name],
                },
            }),
            Some(ToolChoice::Required) => json!({
                "functionCallingConfig": {"mode": "ANY"},
            }),
            Some(ToolChoice::Auto) | None => json!({
                "functionCallingConfig": {"mode": "AUTO"},
            }),
        };
    }
    body
}

/// JSON Schema keywords Gemini's function-calling schema rejects. MCP tool
/// `inputSchema` is full JSON Schema, but Gemini accepts only the OpenAPI
/// subset and 400s on standard keywords like `$schema` or `additionalProperties`.
const GEMINI_UNSUPPORTED_SCHEMA_KEYS: &[&str] = &[
    "$schema",
    "$id",
    "$ref",
    "$comment",
    "$defs",
    "definitions",
    "additionalProperties",
    "unevaluatedProperties",
    "patternProperties",
    "additionalItems",
    "exclusiveMinimum",
    "exclusiveMaximum",
];

/// Recursively drop the keywords Gemini does not accept so a standard JSON
/// Schema is reduced to the subset Gemini's function calling understands.
/// Without this, a tool whose `inputSchema` carries `$schema`,
/// `additionalProperties`, or `exclusiveMinimum` makes the whole request 400.
fn gemini_schema(value: &Value) -> Value {
    match value {
        Value::Object(map) => Value::Object(
            map.iter()
                .filter(|(k, _)| !GEMINI_UNSUPPORTED_SCHEMA_KEYS.contains(&k.as_str()))
                .map(|(k, v)| (k.clone(), gemini_schema(v)))
                .collect(),
        ),
        Value::Array(items) => Value::Array(items.iter().map(gemini_schema).collect()),
        other => other.clone(),
    }
}

async fn handle_response(resp: reqwest::Response) -> Result<Value, ProviderError> {
    let status = resp.status();
    let retry_after = parse_retry_after(resp.headers().get("retry-after"));
    if status.is_success() {
        return resp
            .json::<Value>()
            .await
            .map_err(|e| ProviderError::MalformedResponse(e.to_string()));
    }
    let body = resp.text().await.unwrap_or_default();
    match status.as_u16() {
        401 | 403 => Err(ProviderError::AuthFailure(body)),
        429 => Err(ProviderError::RateLimit { retry_after }),
        400 | 404 | 422 => Err(ProviderError::InvalidRequest(body)),
        500..=599 => Err(ProviderError::ProviderInternal(body)),
        _ => Err(ProviderError::ProviderInternal(format!(
            "unexpected status {status}: {body}"
        ))),
    }
}

fn parse_retry_after(header: Option<&reqwest::header::HeaderValue>) -> Option<Duration> {
    header
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.parse::<u64>().ok())
        .map(Duration::from_secs)
}

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

fn parse_generate_response(value: Value, requested_model: &str) -> Result<Response, ProviderError> {
    let model = value
        .get("modelVersion")
        .and_then(|v| v.as_str())
        .unwrap_or(requested_model)
        .to_owned();
    let candidate = value
        .get("candidates")
        .and_then(|v| v.as_array())
        .and_then(|a| a.first())
        .ok_or_else(|| ProviderError::MalformedResponse("missing candidates[0]".into()))?;
    let finish = candidate
        .get("finishReason")
        .and_then(|v| v.as_str())
        .unwrap_or("STOP");
    let usage_obj: GoogleUsage = serde_json::from_value(
        value
            .get("usageMetadata")
            .cloned()
            .unwrap_or_else(|| Value::Object(Default::default())),
    )
    .map_err(|e| ProviderError::MalformedResponse(format!("usage parse: {e}")))?;
    let usage = Usage {
        input_tokens: usage_obj.prompt_token_count,
        output_tokens: usage_obj.candidates_token_count,
        cache_read_tokens: None,
        cache_write_tokens: None,
    };
    let content = extract_content(candidate)?;
    Ok(Response {
        content,
        finish_reason: map_finish_reason(finish),
        usage,
        model,
    })
}

fn extract_content(candidate: &Value) -> Result<ResponseContent, ProviderError> {
    let parts = candidate
        .pointer("/content/parts")
        .and_then(|v| v.as_array())
        .ok_or_else(|| ProviderError::MalformedResponse("missing content.parts".into()))?;
    let mut text_parts: Vec<String> = Vec::new();
    let mut tool_calls: Vec<ToolCall> = Vec::new();
    for part in parts {
        if let Some(t) = part.get("text").and_then(|v| v.as_str()) {
            text_parts.push(t.to_owned());
            continue;
        }
        if let Some(call) = part.get("functionCall") {
            let name = call
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_owned();
            let arguments = call
                .get("args")
                .cloned()
                .unwrap_or(Value::Object(Default::default()));
            tool_calls.push(ToolCall { name, arguments });
        }
    }
    if !tool_calls.is_empty() {
        return Ok(ResponseContent::ToolUse { calls: tool_calls });
    }
    Ok(ResponseContent::Text {
        text: text_parts.join(""),
    })
}

fn map_finish_reason(s: &str) -> FinishReason {
    match s {
        "STOP" => FinishReason::Stop,
        "MAX_TOKENS" => FinishReason::MaxTokens,
        // Gemini does not emit a dedicated tool-use finish reason; a
        // candidate containing a functionCall part is enough signal. The
        // protocol returns `STOP` in that case, so we treat unknown
        // strings as `Other` and leave tool-use detection to
        // `extract_content`.
        _ => FinishReason::Other,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::eval::provider::ToolSpec;
    use serde_json::json;
    use wiremock::matchers::{body_partial_json, header, method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    fn provider_for(server: &MockServer) -> GoogleProvider {
        GoogleProvider::new(GoogleConfig {
            api_key: "test-key".into(),
            base_url: Some(server.uri()),
            request_timeout: Duration::from_secs(5),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap()
    }

    #[tokio::test]
    async fn happy_path_generate_content_returns_text() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-1.5-flash:generateContent"))
            .and(header("x-goog-api-key", "test-key"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "modelVersion": "gemini-1.5-flash",
                "candidates": [{
                    "content": {
                        "role": "model",
                        "parts": [{"text": "hi there"}],
                    },
                    "finishReason": "STOP",
                }],
                "usageMetadata": {
                    "promptTokenCount": 12,
                    "candidatesTokenCount": 5,
                },
            })))
            .mount(&server)
            .await;

        let provider = provider_for(&server);
        let messages = vec![Message {
            role: Role::User,
            content: "say hi".into(),
        }];
        let options = CompletionOptions {
            model: "gemini-1.5-flash".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&messages, &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "hi there"),
            other => panic!("expected text content, got {other:?}"),
        }
        assert_eq!(response.usage.input_tokens, 12);
        assert_eq!(response.usage.output_tokens, 5);
        assert_eq!(response.finish_reason, FinishReason::Stop);
    }

    #[tokio::test]
    async fn json_mode_sets_response_mime_type() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-2.0-flash:generateContent"))
            .and(body_partial_json(json!({
                "generationConfig": {"responseMimeType": "application/json"},
            })))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "modelVersion": "gemini-2.0-flash",
                "candidates": [{
                    "content": {
                        "role": "model",
                        "parts": [{"text": "{\"ok\":true}"}],
                    },
                    "finishReason": "STOP",
                }],
                "usageMetadata": {
                    "promptTokenCount": 4,
                    "candidatesTokenCount": 4,
                },
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "gemini-2.0-flash".into(),
            json_mode: true,
            ..Default::default()
        };
        let response = provider.chat_completion(&[], &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert!(text.contains("\"ok\"")),
            other => panic!("expected text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn function_call_response_is_parsed_as_tool_calls() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-1.5-pro:generateContent"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "modelVersion": "gemini-1.5-pro",
                "candidates": [{
                    "content": {
                        "role": "model",
                        "parts": [{
                            "functionCall": {
                                "name": "report",
                                "args": {"verdict": "pass"},
                            },
                        }],
                    },
                    "finishReason": "STOP",
                }],
                "usageMetadata": {"promptTokenCount": 30, "candidatesTokenCount": 9},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "gemini-1.5-pro".into(),
            tools: vec![ToolSpec {
                name: "report".into(),
                description: "Return a verdict".into(),
                input_schema: json!({"type": "object"}),
            }],
            tool_choice: Some(ToolChoice::Tool("report".into())),
            ..Default::default()
        };
        let response = provider.chat_completion(&[], &options).await.unwrap();
        match response.content {
            ResponseContent::ToolUse { calls } => {
                assert_eq!(calls.len(), 1);
                assert_eq!(calls[0].name, "report");
                assert_eq!(calls[0].arguments, json!({"verdict": "pass"}));
            }
            other => panic!("expected tool use, got {other:?}"),
        }
    }

    #[test]
    fn forced_tool_call_serializes_function_declarations_and_tool_config() {
        // Google advertises tool_use; a forced tool choice serializes the
        // functionDeclarations array and sets toolConfig to ANY mode with
        // the named function allowed.
        let options = CompletionOptions {
            model: "gemini-1.5-pro".into(),
            tools: vec![ToolSpec {
                name: "report".into(),
                description: "Return a verdict".into(),
                input_schema: json!({"type": "object"}),
            }],
            tool_choice: Some(ToolChoice::Tool("report".into())),
            ..Default::default()
        };
        let body = build_generate_body(&[], &options);
        assert_eq!(
            body.pointer("/tools/0/functionDeclarations/0/name")
                .and_then(|v| v.as_str()),
            Some("report"),
            "function declarations must be serialized into the request body"
        );
        assert_eq!(
            body.pointer("/toolConfig/functionCallingConfig/mode")
                .and_then(|v| v.as_str()),
            Some("ANY")
        );
        assert_eq!(
            body.pointer("/toolConfig/functionCallingConfig/allowedFunctionNames/0")
                .and_then(|v| v.as_str()),
            Some("report")
        );
    }

    #[test]
    fn tool_schema_strips_keywords_gemini_rejects() {
        // MCP inputSchema is full JSON Schema. Gemini 400s on `$schema`,
        // `additionalProperties`, and `exclusiveMinimum`, so the request
        // builder must drop them (recursively) before sending.
        let options = CompletionOptions {
            model: "gemini-2.5-pro".into(),
            tools: vec![ToolSpec {
                name: "get_weather".into(),
                description: "Weather by city".into(),
                input_schema: json!({
                    "$schema": "https://json-schema.org/draft/2020-12/schema",
                    "type": "object",
                    "additionalProperties": false,
                    "properties": {
                        "city": {"type": "string"},
                        "days": {"type": "integer", "exclusiveMinimum": 0}
                    },
                    "required": ["city"]
                }),
            }],
            ..Default::default()
        };
        let body = build_generate_body(&[], &options);
        let params = body
            .pointer("/tools/0/functionDeclarations/0/parameters")
            .expect("parameters serialized");
        assert!(params.get("$schema").is_none(), "$schema must be stripped");
        assert!(
            params.get("additionalProperties").is_none(),
            "additionalProperties must be stripped"
        );
        assert!(
            params
                .pointer("/properties/days/exclusiveMinimum")
                .is_none(),
            "nested exclusiveMinimum must be stripped"
        );
        // Supported fields survive.
        assert_eq!(
            params
                .pointer("/properties/city/type")
                .and_then(|v| v.as_str()),
            Some("string")
        );
        assert_eq!(
            params.pointer("/required/0").and_then(|v| v.as_str()),
            Some("city")
        );
    }

    #[tokio::test]
    async fn count_tokens_uses_dedicated_endpoint() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-1.5-flash:countTokens"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({"totalTokens": 17})))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let count = provider
            .count_tokens("hello world", "gemini-1.5-flash")
            .await
            .unwrap();
        assert_eq!(count, 17);
    }

    #[tokio::test]
    async fn rate_limit_with_retry_after_eventually_succeeds() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-1.5-flash:generateContent"))
            .respond_with(
                ResponseTemplate::new(429)
                    .insert_header("retry-after", "0")
                    .set_body_string("rate limited"),
            )
            .up_to_n_times(2)
            .mount(&server)
            .await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-1.5-flash:generateContent"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "modelVersion": "gemini-1.5-flash",
                "candidates": [{
                    "content": {"role": "model", "parts": [{"text": "ok"}]},
                    "finishReason": "STOP",
                }],
                "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1},
            })))
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "gemini-1.5-flash".into(),
            ..Default::default()
        };
        let response = provider.chat_completion(&[], &options).await.unwrap();
        match response.content {
            ResponseContent::Text { text } => assert_eq!(text, "ok"),
            other => panic!("expected text, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn auth_failure_does_not_retry() {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/models/gemini-1.5-flash:generateContent"))
            .respond_with(ResponseTemplate::new(401).set_body_string("bad key"))
            .expect(1)
            .mount(&server)
            .await;
        let provider = provider_for(&server);
        let options = CompletionOptions {
            model: "gemini-1.5-flash".into(),
            ..Default::default()
        };
        let err = provider.chat_completion(&[], &options).await.unwrap_err();
        assert!(matches!(err, ProviderError::AuthFailure(_)));
    }

    #[test]
    fn estimate_cost_uses_flash_rates_for_default_model() {
        let provider = GoogleProvider::new(GoogleConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            request_timeout: Duration::from_secs(1),
            proxy: crate::network::ProxyConfig::default(),
        })
        .unwrap();
        let usage = Usage {
            input_tokens: 1_000_000,
            output_tokens: 1_000_000,
            ..Default::default()
        };
        let cost = provider.estimate_cost(&usage, "gemini-1.5-flash");
        // 1M input @ $0.075 + 1M output @ $0.30 = $0.375.
        assert!((cost.total() - 0.375).abs() < 1e-6);
    }

    #[test]
    fn estimate_cost_uses_pro_rates_for_pro_model() {
        let provider = GoogleProvider::new(GoogleConfig {
            api_key: "k".into(),
            base_url: Some("http://localhost".into()),
            ..Default::default()
        })
        .unwrap();
        let usage = Usage {
            input_tokens: 1_000_000,
            output_tokens: 1_000_000,
            ..Default::default()
        };
        let cost = provider.estimate_cost(&usage, "gemini-2.0-pro");
        // 1M input @ $1.25 + 1M output @ $5 = $6.25.
        assert!((cost.total() - 6.25).abs() < 1e-6);
    }
}
