aither-gemini 0.21.0

use aither_core::llm::{Classify, RetryPolicy};
use std::time::Duration;

use aither_core::time::sleep;

use serde::{Deserialize, Serialize};
use zenwave::{Client, client, header};

use crate::{
    config::{AuthMode, GeminiConfig, USER_AGENT},
    error::GeminiError,
    types::{
        EmbedContentRequest, EmbedContentResponse, GenerateContentRequest, GenerateContentResponse,
    },
};

/// Response from GET /models/{model} endpoint.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelInfo {
    /// Model name (e.g., "models/gemini-2.0-flash")
    #[allow(dead_code)]
    pub name: String,
    /// Maximum input tokens the model can accept.
    pub input_token_limit: u32,
    /// Maximum output tokens the model can generate.
    #[allow(dead_code)]
    pub output_token_limit: u32,
}

pub async fn call_generate(
    cfg: &GeminiConfig,
    model: &str,
    request: GenerateContentRequest,
) -> Result<GenerateContentResponse, GeminiError> {
    post_json(cfg, cfg.model_endpoint(model, "generateContent"), &request).await
}

/// How a request that failed with a retryable failure is retried: twice,
/// after 2 s and then 4 s, or after the delay the upstream requested.
pub const RETRY_POLICY: RetryPolicy = RetryPolicy {
    max_retries: 2,
    initial_delay: Duration::from_secs(2),
    max_delay: Duration::from_secs(60),
    multiplier: 2,
};

/// Stream generation responses using SSE with retry logic.
pub async fn stream_generate(
    cfg: &GeminiConfig,
    model: &str,
    request: GenerateContentRequest,
) -> Result<
    impl futures_core::Stream<Item = Result<GenerateContentResponse, GeminiError>>,
    GeminiError,
> {
    let mut endpoint = cfg.model_endpoint(model, "streamGenerateContent");
    let separator = if endpoint.contains('?') { '&' } else { '?' };
    endpoint.push(separator);
    endpoint.push_str("alt=sse");
    let debug = std::env::var("AITHER_GEMINI_DEBUG").as_deref() == Ok("1");

    let mut retries = 0u32;
    let sse_stream = loop {
        let mut backend = client();
        let mut builder = backend
            .post(endpoint.clone())
            .map_err(GeminiError::from_http)?
            .header(header::USER_AGENT.as_str(), USER_AGENT)
            .map_err(GeminiError::from_http)?;

        if cfg.auth == AuthMode::Header {
            builder = builder
                .header("x-goog-api-key", cfg.api_key.clone())
                .map_err(GeminiError::from_http)?;
        }

        let builder = builder
            .json_body(&request)
            .map_err(GeminiError::from_http)?;

        match builder.sse().await {
            Ok(stream) => break stream,
            Err(e) => {
                let err = GeminiError::from_http(e);
                let failure = err.failure();
                if RETRY_POLICY.allows(retries, &failure) {
                    let delay = RETRY_POLICY.delay(retries, &failure);
                    retries += 1;
                    if debug {
                        tracing::info!(
                            "Gemini streaming error (retry {}/{}), retrying in {:?}: {}",
                            retries,
                            RETRY_POLICY.max_retries,
                            delay,
                            err
                        );
                    }
                    sleep(delay).await;
                    continue;
                }
                return Err(err);
            }
        }
    };

    Ok(futures_lite::stream::unfold(
        sse_stream,
        |mut stream| async move {
            use futures_lite::StreamExt;
            loop {
                match stream.next().await {
                    None => return None,
                    Some(result) => {
                        let event = match result {
                            Ok(e) => e,
                            Err(e) => {
                                return Some((Err(GeminiError::Stream(e)), stream));
                            }
                        };
                        let data = event.text_data();
                        if data.is_empty() || data == "[DONE]" {
                            continue;
                        }
                        match serde_json::from_str::<GenerateContentResponse>(data) {
                            Ok(response) => return Some((Ok(response), stream)),
                            Err(e) => {
                                tracing::debug!("SSE parse error: {} for data: {}", e, data);
                            }
                        }
                    }
                }
            }
        },
    ))
}

/// Fetch model info including context window size.
pub async fn get_model_info(cfg: &GeminiConfig, model: &str) -> Result<ModelInfo, GeminiError> {
    let model = crate::config::sanitize_model(model);
    get_json(cfg, cfg.endpoint(&model)).await
}

pub async fn embed_content(
    cfg: &GeminiConfig,
    request: EmbedContentRequest,
) -> Result<EmbedContentResponse, GeminiError> {
    post_json(
        cfg,
        cfg.model_endpoint(&cfg.embedding_model, "embedContent"),
        &request,
    )
    .await
}

#[allow(clippy::future_not_send)]
async fn get_json<T: for<'de> serde::Deserialize<'de>>(
    cfg: &GeminiConfig,
    endpoint: String,
) -> Result<T, GeminiError> {
    let mut backend = client();
    let mut builder = backend
        .get(endpoint)
        .map_err(GeminiError::from_http)?
        .header(header::USER_AGENT.as_str(), USER_AGENT)
        .map_err(GeminiError::from_http)?;
    if cfg.auth == AuthMode::Header {
        builder = builder
            .header("x-goog-api-key", cfg.api_key.clone())
            .map_err(GeminiError::from_http)?;
    }
    builder.json().await.map_err(GeminiError::from_http)
}

#[allow(clippy::future_not_send)]
async fn post_json<T: for<'de> serde::Deserialize<'de> + serde::Serialize, S: Serialize>(
    cfg: &GeminiConfig,
    endpoint: String,
    body: &S,
) -> Result<T, GeminiError> {
    let debug = std::env::var("AITHER_GEMINI_DEBUG").as_deref() == Ok("1");
    if debug {
        if let Ok(json) = serde_json::to_string_pretty(body) {
            tracing::debug!("Gemini request to {endpoint}:\n{json}");
        }
        if cfg.auth == AuthMode::Query {
            tracing::debug!("Gemini using query auth (key redacted)");
        }
    }

    let mut retries = 0u32;
    loop {
        let mut backend = client();
        let mut builder = backend
            .post(endpoint.clone())
            .map_err(GeminiError::from_http)?
            .header(header::USER_AGENT.as_str(), USER_AGENT)
            .map_err(GeminiError::from_http)?;
        if cfg.auth == AuthMode::Header {
            builder = builder
                .header("x-goog-api-key", cfg.api_key.clone())
                .map_err(GeminiError::from_http)?;
        }
        let builder = builder.json_body(body).map_err(GeminiError::from_http)?;

        match builder.json().await {
            Ok(res) => {
                if debug && let Ok(json) = serde_json::to_string_pretty(&res) {
                    tracing::debug!("Gemini response from {endpoint}:\n{json}");
                }
                return Ok(res);
            }
            Err(e) => {
                let err = GeminiError::from_http(e);
                let failure = err.failure();
                if RETRY_POLICY.allows(retries, &failure) {
                    let delay = RETRY_POLICY.delay(retries, &failure);
                    retries += 1;
                    if debug {
                        tracing::info!(
                            "Gemini POST error (retry {}/{}), retrying in {:?}: {}",
                            retries,
                            RETRY_POLICY.max_retries,
                            delay,
                            err
                        );
                    }
                    sleep(delay).await;
                    continue;
                }
                return Err(err);
            }
        }
    }
}