//! Gemini-native endpoint — zero-translation tool call passthrough.
//!
//! Routes:
//!   POST /v1beta/models/{model}:generateContent      (sync, official Gemini API path)
//!   POST /v1beta/models/{model}:streamGenerateContent (stream, official Gemini API path)
//!
//! Accepts tools in Gemini `functionDeclarations` format directly,
//! returns `functionCall` in Gemini format directly.
//! No OpenAI ↔ Gemini format conversion.

use axum::{
    extract::State,
    http::StatusCode,
    response::{sse::Event, IntoResponse, Json, Sse},
};
use rand::Rng;
use std::sync::Arc;
use tracing::{debug, info, warn};

use super::models::{lookup_model, DEFAULT_MODEL, MODELS};
use super::polling::{
    extract_response_text, extract_thinking_content, is_response_done, poll_for_response,
};
use super::util::{err_response, upstream_err_response};
use super::AppState;
use crate::mitm::store::PendingToolResult;

/// Gemini-native request format.
#[derive(serde::Deserialize)]
pub(crate) struct GeminiRequest {
    pub model: Option<String>,
    /// User input text (our custom format).
    #[serde(default)]
    pub input: Option<serde_json::Value>,
    /// Official Gemini API format: [{"role": "user", "parts": [{"text": "..."}]}]
    #[serde(default)]
    pub contents: Option<Vec<serde_json::Value>>,
    /// Shorthand: single text message (alias for simple requests).
    #[serde(default)]
    pub message: Option<String>,
    /// Gemini-native tools: [{"functionDeclarations": [...]}]
    #[serde(default)]
    pub tools: Option<Vec<serde_json::Value>>,
    /// Gemini-native toolConfig: {"functionCallingConfig": {"mode": "AUTO"}}
    #[serde(default)]
    pub tool_config: Option<serde_json::Value>,
    /// Session/conversation ID.
    #[serde(default)]
    #[allow(dead_code)]
    pub conversation: Option<serde_json::Value>,
    #[serde(default = "default_timeout")]
    pub timeout: u64,
    #[serde(default)]
    pub stream: bool,
    /// Tool results in Gemini format: [{"functionResponse": {"name": "...", "response": {...}}}]
    #[serde(default)]
    pub tool_results: Option<Vec<serde_json::Value>>,
    /// Sampling temperature.
    #[serde(default)]
    pub temperature: Option<f64>,
    /// Nucleus sampling (topP).
    #[serde(default, alias = "topP")]
    pub top_p: Option<f64>,
    /// Top-K sampling.
    #[serde(default, alias = "topK")]
    pub top_k: Option<u32>,
    /// Max output tokens.
    #[serde(default, alias = "maxOutputTokens")]
    pub max_output_tokens: Option<u64>,
    /// Stop sequences.
    #[serde(default, alias = "stopSequences")]
    pub stop_sequences: Option<Vec<String>>,
    /// Thinking level for Gemini 3 models: "minimal", "low", "medium", "high".
    /// Maps directly to thinkingConfig.thinkingLevel.
    #[serde(default, alias = "thinkingLevel")]
    pub thinking_level: Option<String>,
    /// Enable Google Search grounding. See Gemini API docs.
    /// When true, injects {"googleSearch": {}} into tools via MITM.
    #[serde(default, alias = "googleSearch")]
    pub google_search: bool,
    /// Response MIME type for structured output, e.g. "application/json".
    /// Injected as generationConfig.responseMimeType via MITM.
    #[serde(default, alias = "responseMimeType")]
    pub response_mime_type: Option<String>,
    /// Response schema for structured output (JSON Schema object).
    /// Injected as generationConfig.responseSchema via MITM.
    #[serde(default, alias = "responseSchema")]
    pub response_schema: Option<serde_json::Value>,
}

use super::util::default_timeout;

/// Build Gemini-format usageMetadata from MITM store.
async fn build_usage_metadata(
    store: &crate::mitm::store::MitmStore,
    cascade_id: &str,
) -> serde_json::Value {
    let usage = store
        .take_usage(cascade_id)
        .await
        .or(store.take_usage("_latest").await);
    if let Some(usage) = usage {
        serde_json::json!({
            "promptTokenCount": usage.input_tokens,
            "candidatesTokenCount": usage.output_tokens,
            "totalTokenCount": usage.input_tokens + usage.output_tokens,
            "thoughtsTokenCount": usage.thinking_output_tokens,
            "cachedContentTokenCount": usage.cache_read_input_tokens,
        })
    } else {
        serde_json::json!({
            "promptTokenCount": 0,
            "candidatesTokenCount": 0,
            "totalTokenCount": 0,
        })
    }
}

/// POST /v1beta/*path — handles both :generateContent and :streamGenerateContent
///
/// Parses paths like:
///   /v1beta/models/gemini-3-flash:generateContent
///   /v1beta/models/gemini-3-flash:streamGenerateContent
pub(crate) async fn handle_gemini_v1beta(
    State(state): State<Arc<AppState>>,
    axum::extract::Path(path): axum::extract::Path<String>,
    Json(mut body): Json<GeminiRequest>,
) -> axum::response::Response {
    // Expected path: "models/{model}:{action}"
    let path = path.trim_start_matches('/');
    let model_action = path.strip_prefix("models/").unwrap_or(path);

    // Split on the LAST colon — model name may contain hyphens/numbers but not colons
    if let Some(colon_pos) = model_action.rfind(':') {
        let model = &model_action[..colon_pos];
        let action = &model_action[colon_pos + 1..];

        body.model = Some(model.to_string());
        match action {
            "generateContent" => body.stream = false,
            "streamGenerateContent" => body.stream = true,
            _ => {
                return err_response(
                    StatusCode::BAD_REQUEST,
                    format!(
                        "Unknown action: {action}. Use :generateContent or :streamGenerateContent"
                    ),
                    "invalid_request_error",
                );
            }
        }
    } else {
        return err_response(
            StatusCode::BAD_REQUEST,
            format!(
                "Invalid path: /v1beta/{path}. Expected /v1beta/models/{{model}}:generateContent"
            ),
            "invalid_request_error",
        );
    }

    handle_gemini_inner(state, body).await
}

async fn handle_gemini_inner(
    state: Arc<AppState>,
    body: GeminiRequest,
) -> axum::response::Response {
    info!(
        "POST /v1beta/gemini model={} stream={}",
        body.model.as_deref().unwrap_or(DEFAULT_MODEL),
        body.stream
    );

    let model_name = body.model.as_deref().unwrap_or(DEFAULT_MODEL);
    let model = match lookup_model(model_name) {
        Some(m) => m,
        None => {
            let names: Vec<&str> = MODELS.iter().map(|m| m.name).collect();
            return err_response(
                StatusCode::BAD_REQUEST,
                format!("Unknown model: {model_name}. Available: {names:?}"),
                "invalid_request_error",
            );
        }
    };

    let token = state.backend.oauth_token().await;
    if token.is_empty() {
        return err_response(
            StatusCode::UNAUTHORIZED,
            "No OAuth token. POST to /v1/token or set ZEROGRAVITY_TOKEN env var.".into(),
            "authentication_error",
        );
    }

    // Extract user text and optional image.
    // Priority: contents (official Gemini API) > input (our format) > message (shorthand)
    let mut image: Option<crate::proto::ImageData> = None;
    let user_text = if let Some(ref contents) = body.contents {
        // Official Gemini API format: [{"role": "user", "parts": [{"text": "..."}]}]
        // Extract text from the last user message.
        let mut text_parts: Vec<String> = Vec::new();
        for content in contents.iter().rev() {
            let role = content
                .get("role")
                .and_then(|r| r.as_str())
                .unwrap_or("user");
            if role != "user" {
                continue;
            }
            if let Some(parts) = content.get("parts").and_then(|p| p.as_array()) {
                for part in parts {
                    if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
                        text_parts.push(text.to_string());
                    }
                    // Handle inlineData image
                    if image.is_none() {
                        if let Some(inline) = part.get("inlineData") {
                            if let (Some(mime), Some(b64)) =
                                (inline["mimeType"].as_str(), inline["data"].as_str())
                            {
                                if let Some(img) = super::util::parse_data_uri(&format!(
                                    "data:{mime};base64,{b64}"
                                )) {
                                    image = Some(img);
                                }
                            }
                        }
                    }
                }
            }
            if !text_parts.is_empty() {
                break;
            }
        }
        if text_parts.is_empty() {
            return err_response(
                StatusCode::BAD_REQUEST,
                "No text found in contents array".to_string(),
                "invalid_request_error",
            );
        }
        text_parts.join("\n")
    } else if let Some(ref input) = body.input {
        // Our custom format: input as string or array
        match input {
            serde_json::Value::String(s) => s.clone(),
            serde_json::Value::Array(arr) => {
                let mut parts: Vec<String> = Vec::new();
                for item in arr {
                    match item {
                        serde_json::Value::String(s) => parts.push(s.clone()),
                        serde_json::Value::Object(obj) => {
                            if let Some(text) = obj.get("text").and_then(|v| v.as_str()) {
                                parts.push(text.to_string());
                            }
                            if image.is_none() {
                                if let Some(inline) = obj.get("inlineData") {
                                    if let (Some(mime), Some(b64)) =
                                        (inline["mimeType"].as_str(), inline["data"].as_str())
                                    {
                                        if let Some(img) = super::util::parse_data_uri(&format!(
                                            "data:{mime};base64,{b64}"
                                        )) {
                                            image = Some(img);
                                        }
                                    }
                                }
                                if let Some(img) = super::util::extract_image_from_content(item) {
                                    image = Some(img);
                                }
                            }
                        }
                        _ => {}
                    }
                }
                if parts.is_empty() {
                    return err_response(
                        StatusCode::BAD_REQUEST,
                        "Gemini input array contains no text parts".to_string(),
                        "invalid_request_error",
                    );
                }
                parts.join("\n")
            }
            _ => {
                return err_response(
                    StatusCode::BAD_REQUEST,
                    "Gemini input must be a string or array of text parts".to_string(),
                    "invalid_request_error",
                );
            }
        }
    } else if let Some(ref msg) = body.message {
        msg.clone()
    } else {
        return err_response(
            StatusCode::BAD_REQUEST,
            "Request must include 'contents' (Gemini API), 'input', or 'message'".to_string(),
            "invalid_request_error",
        );
    };

    // ── Build per-request state locally ──────────────────────────────────

    // Tools (already in Gemini format)
    let tools = body.tools.as_ref().and_then(|t| {
        if t.is_empty() {
            None
        } else {
            info!(count = t.len(), "Gemini-native tools for MITM injection");
            Some(t.clone())
        }
    });
    let tool_config = body.tool_config.clone();

    // Tool results → collect (ToolRound built after cascade_id is known)
    let mut pending_tool_results: Vec<PendingToolResult> = Vec::new();
    if let Some(ref results) = body.tool_results {
        for r in results {
            if let Some(fr) = r.get("functionResponse") {
                let name = fr["name"].as_str().unwrap_or("unknown").to_string();
                let response = fr.get("response").cloned().unwrap_or(serde_json::json!({}));
                pending_tool_results.push(PendingToolResult {
                    name,
                    result: response,
                });
            }
        }
        info!(
            count = results.len(),
            "Gemini-native tool results (will build tool round after cascade_id)"
        );
    }

    // Generation parameters
    use crate::mitm::store::GenerationParams;
    let gp = GenerationParams {
        temperature: body.temperature,
        top_p: body.top_p,
        top_k: body.top_k,
        max_output_tokens: body.max_output_tokens,
        stop_sequences: body.stop_sequences.clone(),
        frequency_penalty: None,
        presence_penalty: None,
        reasoning_effort: body.thinking_level.clone(),
        response_mime_type: body.response_mime_type.clone(),
        response_schema: body.response_schema.clone(),
        google_search: body.google_search,
    };
    let generation_params = if gp.temperature.is_some()
        || gp.top_p.is_some()
        || gp.top_k.is_some()
        || gp.max_output_tokens.is_some()
        || gp.stop_sequences.is_some()
        || gp.reasoning_effort.is_some()
        || gp.response_mime_type.is_some()
        || gp.response_schema.is_some()
        || gp.google_search
    {
        Some(gp)
    } else {
        None
    };

    // Always create a new cascade for every request
    let cascade_id = match state.backend.create_cascade().await {
        Ok(cid) => cid,
        Err(e) => {
            return err_response(
                StatusCode::BAD_GATEWAY,
                format!("StartCascade failed: {e}"),
                "server_error",
            );
        }
    };

    // Image for MITM injection
    let pending_image = image.as_ref().map(|img| {
        use base64::Engine;
        crate::mitm::store::PendingImage {
            base64_data: base64::engine::general_purpose::STANDARD.encode(&img.data),
            mime_type: img.mime_type.clone(),
        }
    });

    // Build event channel — always created for MITM response path
    let (tx, rx) = tokio::sync::mpsc::channel(64);
    let (mitm_rx, event_tx) = (Some(rx), tx);

    // Build tool rounds now that cascade_id is known
    let mut tool_rounds: Vec<crate::mitm::store::ToolRound> = Vec::new();
    if !pending_tool_results.is_empty() {
        let last_calls = state
            .mitm_store
            .take_function_calls(&cascade_id)
            .await
            .unwrap_or_default();
        tool_rounds.push(crate::mitm::store::ToolRound {
            calls: last_calls,
            results: pending_tool_results.clone(),
        });
    }

    // Start debug trace
    let trace = state
        .trace
        .start(&cascade_id, "POST gemini", model_name, body.stream);
    if let Some(ref t) = trace {
        t.set_client_request(crate::trace::ClientRequestSummary {
            message_count: 1,
            tool_count: body.tools.as_ref().map_or(0, |t| t.len()),
            tool_round_count: tool_rounds.len(),
            user_text_len: user_text.len(),
            user_text_preview: user_text.chars().take(200).collect(),
            system_prompt: false,
            has_image: image.is_some(),
        })
        .await;
        t.start_turn().await;
    }

    let mitm_gate = std::sync::Arc::new(tokio::sync::Notify::new());
    let mitm_gate_clone = mitm_gate.clone();
    state
        .mitm_store
        .register_request(crate::mitm::store::RequestContext {
            cascade_id: cascade_id.clone(),
            pending_user_text: user_text.clone(),
            event_channel: event_tx,
            generation_params,
            pending_image,
            tools,
            tool_config,
            pending_tool_results,
            tool_rounds,
            last_function_calls: Vec::new(),
            call_id_to_name: std::collections::HashMap::new(),
            created_at: std::time::Instant::now(),
            gate: mitm_gate_clone,
            trace_handle: trace.clone(),
            trace_turn: 0,
        })
        .await;

    // Send REAL user text to LS (no more dummy ".")
    match state
        .backend
        .send_message_with_image(
            &cascade_id,
            &format!(".<cid:{}>", cascade_id),
            model.model_enum,
            image.as_ref(),
        )
        .await
    {
        Ok((200, _)) => {
            let bg = Arc::clone(&state.backend);
            let cid = cascade_id.clone();
            tokio::spawn(async move {
                let _ = bg.update_annotations(&cid).await;
            });
        }
        Ok((status, _)) => {
            state.mitm_store.remove_request(&cascade_id).await;
            return err_response(
                StatusCode::BAD_GATEWAY,
                format!("Antigravity returned {status}"),
                "server_error",
            );
        }
        Err(e) => {
            state.mitm_store.remove_request(&cascade_id).await;
            return err_response(
                StatusCode::BAD_GATEWAY,
                format!("Send message failed: {e}"),
                "server_error",
            );
        }
    }

    // Wait for MITM gate: 5s -> 502 if MITM enabled
    let gate_start = std::time::Instant::now();
    let gate_matched =
        tokio::time::timeout(std::time::Duration::from_secs(5), mitm_gate.notified()).await;
    let gate_wait_ms = gate_start.elapsed().as_millis() as u64;
    if gate_matched.is_err() {
        if state.mitm_enabled {
            state.mitm_store.remove_request(&cascade_id).await;
            if let Some(ref t) = trace {
                t.record_error("MITM gate timeout (5s)".to_string()).await;
                t.finish("mitm_timeout").await;
            }
            return err_response(
                StatusCode::BAD_GATEWAY,
                "MITM proxy did not match request within 5s".to_string(),
                "mitm_timeout",
            );
        }
        warn!(cascade = %cascade_id, "MITM gate timeout (--no-mitm mode)");
    } else {
        debug!(cascade = %cascade_id, gate_wait_ms, "MITM gate signaled -- request matched");
        if let Some(ref t) = trace {
            t.record_mitm_match(0, gate_wait_ms).await;
        }
    }

    // Dispatch to sync or stream
    let model_name = model_name.to_string();
    let timeout = body.timeout;
    if body.stream {
        gemini_stream(state, model_name, cascade_id, timeout, mitm_rx, trace).await
    } else {
        gemini_sync(state, model_name, cascade_id, timeout, mitm_rx, trace).await
    }
}

// ─── Sync ────────────────────────────────────────────────────────────────────

async fn gemini_sync(
    state: Arc<AppState>,
    model_name: String,
    cascade_id: String,
    timeout: u64,
    mitm_rx: Option<tokio::sync::mpsc::Receiver<crate::mitm::store::MitmEvent>>,
    trace: Option<crate::trace::TraceHandle>,
) -> axum::response::Response {
    // Clear stale response and upstream errors (only if no pre-installed channel)
    if mitm_rx.is_none() {
        state.mitm_store.clear_response_async().await;
        state.mitm_store.clear_upstream_error().await;
    }

    // ── MITM bypass: channel-based pipeline when custom tools active ──
    if let Some(mut rx) = mitm_rx {
        let start = std::time::Instant::now();

        let mut acc_text = String::new();
        let mut acc_thinking: Option<String> = None;
        let mut last_usage: Option<crate::mitm::store::ApiUsage> = None;

        while let Some(event) = tokio::time::timeout(
            std::time::Duration::from_secs(timeout.saturating_sub(start.elapsed().as_secs())),
            rx.recv(),
        )
        .await
        .ok()
        .flatten()
        {
            use crate::mitm::store::MitmEvent;
            match event {
                MitmEvent::ThinkingDelta(t) => {
                    acc_thinking = Some(t);
                }
                MitmEvent::TextDelta(t) => {
                    acc_text = t;
                }
                MitmEvent::Usage(u) => {
                    last_usage = Some(u);
                }
                MitmEvent::Grounding(_) => {}
                MitmEvent::FunctionCall(calls) => {
                    let parts: Vec<serde_json::Value> = calls
                        .iter()
                        .map(|fc| {
                            serde_json::json!({
                                "functionCall": {
                                    "name": fc.name,
                                    "args": fc.args,
                                }
                            })
                        })
                        .collect();
                    if let Some(ref t) = trace {
                        let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls
                            .iter()
                            .map(|fc| crate::trace::FunctionCallSummary {
                                name: fc.name.clone(),
                                args_preview: serde_json::to_string(&fc.args)
                                    .unwrap_or_default()
                                    .chars()
                                    .take(200)
                                    .collect(),
                            })
                            .collect();
                        t.record_response(
                            0,
                            crate::trace::ResponseSummary {
                                text_len: 0,
                                thinking_len: acc_thinking.as_ref().map_or(0, |s| s.len()),
                                text_preview: String::new(),
                                finish_reason: Some("STOP".to_string()),
                                function_calls: fc_summaries,
                                grounding: false,
                            },
                        )
                        .await;
                        t.finish("tool_call").await;
                    }
                    state.mitm_store.remove_request(&cascade_id).await;
                    return Json(serde_json::json!({
                        "candidates": [{
                            "content": {
                                "parts": parts,
                                "role": "model",
                            },
                            "finishReason": "STOP",
                        }],
                        "modelVersion": model_name,
                    }))
                    .into_response();
                }
                MitmEvent::ResponseComplete => {
                    if acc_text.is_empty() && acc_thinking.is_none() {
                        // Empty response — continue waiting
                        continue;
                    }
                    if acc_text.is_empty() && acc_thinking.is_some() {
                        // Thinking-only — LS needs to make a follow-up request.
                        // Reinstall channel and unblock gate.
                        let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
                        state.mitm_store.set_channel(&cascade_id, new_tx).await;

                        let _ = state.mitm_store.take_any_function_calls().await;
                        rx = new_rx;
                        debug!(
                            "Gemini sync: thinking-only — new channel for follow-up, thinking_len={}",
                            acc_thinking.as_ref().map(|t| t.len()).unwrap_or(0)
                        );
                        continue;
                    }
                    let mut parts: Vec<serde_json::Value> = Vec::new();
                    if let Some(ref t) = acc_thinking {
                        parts.push(serde_json::json!({"text": t, "thought": true}));
                    }
                    parts.push(serde_json::json!({"text": acc_text}));
                    if let Some(ref t) = trace {
                        t.record_response(
                            0,
                            crate::trace::ResponseSummary {
                                text_len: acc_text.len(),
                                thinking_len: acc_thinking.as_ref().map_or(0, |s| s.len()),
                                text_preview: acc_text.chars().take(200).collect(),
                                finish_reason: Some("STOP".to_string()),
                                function_calls: Vec::new(),
                                grounding: false,
                            },
                        )
                        .await;
                        if let Some(ref u) = last_usage {
                            t.set_usage(crate::trace::TrackedUsage {
                                input_tokens: u.input_tokens,
                                output_tokens: u.output_tokens,
                                thinking_tokens: u.thinking_output_tokens,
                                cache_read: u.cache_read_input_tokens,
                            })
                            .await;
                        }
                        t.finish("completed").await;
                    }
                    state.mitm_store.remove_request(&cascade_id).await;
                    return Json(serde_json::json!({
                        "candidates": [{
                            "content": {
                                "parts": parts,
                                "role": "model",
                            },
                            "finishReason": "STOP",
                        }],
                        "modelVersion": model_name,
                        "usageMetadata": if let Some(ref u) = last_usage {
                            serde_json::json!({
                                "promptTokenCount": u.input_tokens,
                                "candidatesTokenCount": u.output_tokens,
                                "totalTokenCount": u.input_tokens + u.output_tokens,
                                "thoughtsTokenCount": u.thinking_output_tokens,
                                "cachedContentTokenCount": u.cache_read_input_tokens,
                            })
                        } else {
                            build_usage_metadata(&state.mitm_store, &cascade_id).await
                        },
                    }))
                    .into_response();
                }
                MitmEvent::UpstreamError(err) => {
                    if let Some(ref t) = trace {
                        t.record_response(
                            0,
                            crate::trace::ResponseSummary {
                                text_len: acc_text.len(),
                                thinking_len: acc_thinking.as_ref().map_or(0, |s| s.len()),
                                text_preview: acc_text.chars().take(200).collect(),
                                finish_reason: Some("STOP".to_string()),
                                function_calls: Vec::new(),
                                grounding: false,
                            },
                        )
                        .await;
                        if let Some(ref u) = last_usage {
                            t.set_usage(crate::trace::TrackedUsage {
                                input_tokens: u.input_tokens,
                                output_tokens: u.output_tokens,
                                thinking_tokens: u.thinking_output_tokens,
                                cache_read: u.cache_read_input_tokens,
                            })
                            .await;
                        }
                        t.finish("upstream_error").await;
                    }
                    state.mitm_store.remove_request(&cascade_id).await;
                    return upstream_err_response(&err);
                }
            }
        }

        // Timeout
        if let Some(ref t) = trace {
            t.record_error(format!("Timeout: no response after {timeout}s"))
                .await;
            t.finish("timeout").await;
        }
        state.mitm_store.remove_request(&cascade_id).await;
        return (
            axum::http::StatusCode::GATEWAY_TIMEOUT,
            Json(serde_json::json!({
                "error": {
                    "message": format!("Timeout: no response from Google API after {timeout}s"),
                    "type": "upstream_error",
                    "code": 504,
                }
            })),
        )
            .into_response();
    }

    // ── Normal LS path (no custom tools) ──
    let poll_result = poll_for_response(&state, &cascade_id, timeout).await;
    if let Some(ref err) = poll_result.upstream_error {
        return upstream_err_response(err);
    }

    // Check for captured function calls — return in Gemini format
    let captured_tool_calls = state.mitm_store.take_any_function_calls().await;

    if let Some(ref calls) = captured_tool_calls {
        info!(
            count = calls.len(),
            tools = ?calls.iter().map(|c| &c.name).collect::<Vec<_>>(),
            "Returning captured function calls (Gemini format)"
        );

        let parts: Vec<serde_json::Value> = calls
            .iter()
            .map(|fc| {
                serde_json::json!({
                    "functionCall": {
                        "name": fc.name,
                        "args": fc.args,
                    }
                })
            })
            .collect();

        // Record trace
        if let Some(ref t) = trace {
            let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls
                .iter()
                .map(|fc| crate::trace::FunctionCallSummary {
                    name: fc.name.clone(),
                    args_preview: serde_json::to_string(&fc.args)
                        .unwrap_or_default()
                        .chars()
                        .take(200)
                        .collect(),
                })
                .collect();
            t.record_response(
                0,
                crate::trace::ResponseSummary {
                    text_len: 0,
                    thinking_len: 0,
                    text_preview: String::new(),
                    finish_reason: Some("STOP".to_string()),
                    function_calls: fc_summaries,
                    grounding: false,
                },
            )
            .await;
            t.finish("tool_call").await;
        }

        return Json(serde_json::json!({
            "candidates": [{
                "content": {
                    "parts": parts,
                    "role": "model",
                },
                "finishReason": "STOP",
            }],
            "modelVersion": model_name,
            "usageMetadata": build_usage_metadata(&state.mitm_store, &cascade_id).await,
        }))
        .into_response();
    }

    // Normal text response — include thinking if present
    let mut parts: Vec<serde_json::Value> = Vec::new();
    if let Some(ref thinking) = poll_result.thinking {
        parts.push(serde_json::json!({"text": thinking, "thought": true}));
    }
    parts.push(serde_json::json!({"text": poll_result.text}));

    // Record trace
    if let Some(ref t) = trace {
        t.record_response(
            0,
            crate::trace::ResponseSummary {
                text_len: poll_result.text.len(),
                thinking_len: poll_result.thinking.as_ref().map_or(0, |s| s.len()),
                text_preview: poll_result.text.chars().take(200).collect(),
                finish_reason: Some("STOP".to_string()),
                function_calls: Vec::new(),
                grounding: false,
            },
        )
        .await;
        t.finish("completed").await;
    }

    Json(serde_json::json!({
        "candidates": [{
            "content": {
                "parts": parts,
                "role": "model",
            },
            "finishReason": "STOP",
        }],
        "modelVersion": model_name,
        "usageMetadata": build_usage_metadata(&state.mitm_store, &cascade_id).await,
    }))
    .into_response()
}

// ─── Streaming ───────────────────────────────────────────────────────────────

async fn gemini_stream(
    state: Arc<AppState>,
    model_name: String,
    cascade_id: String,
    timeout: u64,
    mitm_rx: Option<tokio::sync::mpsc::Receiver<crate::mitm::store::MitmEvent>>,
    trace: Option<crate::trace::TraceHandle>,
) -> axum::response::Response {
    let stream = async_stream::stream! {
        let start = std::time::Instant::now();
        let mut last_text = String::new();
        let mut last_thinking = String::new();
        let mut last_usage: Option<crate::mitm::store::ApiUsage> = None;

        // Clear stale response (only if no pre-installed channel)
        if mitm_rx.is_none() {
            state.mitm_store.clear_response_async().await;
        }

        // ── MITM bypass path (custom tools active) ──
        // Channel-based pipeline: read events directly from MITM proxy.
        // Channel is pre-installed before send_message to avoid race conditions.
        if let Some(mut rx) = mitm_rx {
            let mut did_unblock_ls = false;

            'channel_loop: while let Some(event) = tokio::time::timeout(
                std::time::Duration::from_secs(timeout.saturating_sub(start.elapsed().as_secs())),
                rx.recv(),
            ).await.ok().flatten() {
                use crate::mitm::store::MitmEvent;
                match event {
                    MitmEvent::ThinkingDelta(full_thinking) => {
                        if full_thinking.len() > last_thinking.len() {
                            let delta = full_thinking[last_thinking.len()..].to_string();
                            last_thinking = full_thinking;

                            yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
                                "candidates": [{
                                    "content": {
                                        "parts": [{"text": delta, "thought": true}],
                                        "role": "model",
                                    },
                                }],
                                "modelVersion": model_name,
                            })).unwrap_or_default()));
                        }
                    }
                    MitmEvent::TextDelta(full_text) => {
                        if full_text.len() > last_text.len() {
                            let delta = full_text[last_text.len()..].to_string();

                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                                "candidates": [{
                                    "content": {
                                        "parts": [{"text": delta}],
                                        "role": "model",
                                    },
                                }],
                                "modelVersion": model_name,
                            })).unwrap_or_default()));
                            last_text = full_text;
                        }
                    }
                    MitmEvent::FunctionCall(calls) => {
                        let parts: Vec<serde_json::Value> = calls
                            .iter()
                            .map(|fc| {
                                serde_json::json!({
                                    "functionCall": {
                                        "name": fc.name,
                                        "args": fc.args,
                                    }
                                })
                            })
                            .collect();

                        let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await;
                        yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
                            "candidates": [{
                                "content": {
                                    "parts": parts,
                                    "role": "model",
                                },
                                "finishReason": "STOP",
                            }],
                            "usageMetadata": usage_meta,
                            "modelVersion": model_name,
                        })).unwrap_or_default()));
                        yield Ok(Event::default().data("[DONE]"));
                        if let Some(ref t) = trace {
                            let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls.iter().map(|fc| crate::trace::FunctionCallSummary {
                                name: fc.name.clone(), args_preview: serde_json::to_string(&fc.args).unwrap_or_default().chars().take(200).collect(),
                            }).collect();
                            t.record_response(0, crate::trace::ResponseSummary {
                                text_len: 0, thinking_len: last_thinking.len(), text_preview: String::new(),
                                finish_reason: Some("STOP".to_string()),
                                function_calls: fc_summaries, grounding: false,
                            }).await;
                            t.finish("tool_call").await;
                        }
                        state.mitm_store.remove_request(&cascade_id).await;
                        return;
                    }
                    MitmEvent::ResponseComplete => {
                        if !last_text.is_empty() {
                            // Final chunk with finishReason + usageMetadata
                            let usage_meta = if let Some(ref u) = last_usage {
                                serde_json::json!({
                                    "promptTokenCount": u.input_tokens,
                                    "candidatesTokenCount": u.output_tokens,
                                    "totalTokenCount": u.input_tokens + u.output_tokens,
                                    "thoughtsTokenCount": u.thinking_output_tokens,
                                    "cachedContentTokenCount": u.cache_read_input_tokens,
                                })
                            } else {
                                build_usage_metadata(&state.mitm_store, &cascade_id).await
                            };
                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                                "candidates": [{
                                    "content": {
                                        "parts": [{"text": ""}],
                                        "role": "model",
                                    },
                                    "finishReason": "STOP",
                                }],
                                "usageMetadata": usage_meta,
                                "modelVersion": model_name,
                            })).unwrap_or_default()));
                            yield Ok(Event::default().data("[DONE]"));
                            if let Some(ref t) = trace {
                                t.record_response(0, crate::trace::ResponseSummary {
                                    text_len: last_text.len(), thinking_len: last_thinking.len(),
                                    text_preview: last_text.chars().take(200).collect(),
                                    finish_reason: Some("STOP".to_string()),
                                    function_calls: Vec::new(), grounding: false,
                                }).await;
                                t.finish("completed").await;
                            }
                            state.mitm_store.remove_request(&cascade_id).await;
                            return;
                        } else if !last_thinking.is_empty() && !did_unblock_ls {
                            // Thinking-only response — LS needs follow-up API calls.
                            // Create a new channel and unblock the gate.
                            did_unblock_ls = true;
                            let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
                            state.mitm_store.set_channel(&cascade_id, new_tx).await;

                            let _ = state.mitm_store.take_any_function_calls().await;
                            rx = new_rx;
                            debug!(
                                "Gemini stream: thinking-only — new channel for follow-up, thinking_len={}",
                                last_thinking.len()
                            );
                            continue 'channel_loop;
                        }
                        // Empty response — continue waiting
                    }
                    MitmEvent::UpstreamError(err) => {
                        let error_msg = super::util::upstream_error_message(&err);
                        let error_type = super::util::upstream_error_type(&err);
                        yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
                            "error": {
                                "message": error_msg,
                                "type": error_type,
                                "code": err.status,
                            }
                        })).unwrap()));
                        yield Ok(Event::default().data("[DONE]"));
                        if let Some(ref t) = trace {
                            t.record_error(format!("Upstream: {}", error_msg)).await;
                            t.finish("upstream_error").await;
                        }
                        state.mitm_store.remove_request(&cascade_id).await;
                        return;
                    }
                    MitmEvent::Usage(u) => { last_usage = Some(u); }
                    MitmEvent::Grounding(_) => {}
                }
            }

            // Timeout or channel closed
            state.mitm_store.remove_request(&cascade_id).await;
            yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({
                "error": {
                    "message": format!("Timeout: no response from Google API after {timeout}s"),
                    "type": "upstream_error",
                    "code": 504,
                }
            })).unwrap()));
            if let Some(ref t) = trace {
                t.record_error(format!("Timeout: {timeout}s")).await;
                t.finish("timeout").await;
            }
            yield Ok(Event::default().data("[DONE]"));
            return;
        }

        while start.elapsed().as_secs() < timeout {

            // ── Normal LS path (no custom tools) ──
            if let Ok((status, data)) = state.backend.get_steps(&cascade_id).await {
                if status == 200 {
                    if let Some(steps) = data["steps"].as_array() {
                        // Stream thinking
                        if let Some(tc) = extract_thinking_content(steps) {
                            if tc.len() > last_thinking.len() {
                                let delta = tc[last_thinking.len()..].to_string();
                                last_thinking = tc;

                                yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                                    "candidates": [{
                                        "content": {
                                            "parts": [{"text": delta, "thought": true}],
                                            "role": "model",
                                        },
                                    }],
                                    "modelVersion": model_name,
                                })).unwrap_or_default()));
                            }
                        }

                        // Stream text
                        let text = extract_response_text(steps);
                        if !text.is_empty() && text.len() > last_text.len() {
                            let delta = if text.starts_with(&*last_text) {
                                &text[last_text.len()..]
                            } else {
                                &text
                            };

                            if !delta.is_empty() {
                                yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                                    "candidates": [{
                                        "content": {
                                            "parts": [{"text": delta}],
                                            "role": "model",
                                        },
                                    }],
                                    "modelVersion": model_name,
                                })).unwrap_or_default()));
                                last_text = text.to_string();
                            }
                        }

                        // Done check
                        if is_response_done(steps) && !last_text.is_empty() {
                            let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await;
                            yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
                                "candidates": [{
                                    "content": {
                                        "parts": [{"text": ""}],
                                        "role": "model",
                                    },
                                    "finishReason": "STOP",
                                }],
                                "usageMetadata": usage_meta,
                                "modelVersion": model_name,
                            })).unwrap_or_default()));
                            yield Ok(Event::default().data("[DONE]"));
                            return;
                        }
                    }
                }
            }

            let poll_ms: u64 = rand::thread_rng().gen_range(200..350);
            tokio::time::sleep(tokio::time::Duration::from_millis(poll_ms)).await;
        }

        // Timeout — emit proper error
        yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
            "error": {
                "message": format!("Timeout: no response from Google API after {timeout}s"),
                "type": "upstream_error",
                "code": 504,
            }
        })).unwrap_or_default()));
        yield Ok(Event::default().data("[DONE]"));
    };

    Sse::new(stream)
        .keep_alive(
            axum::response::sse::KeepAlive::new()
                .interval(std::time::Duration::from_secs(15))
                .text(""),
        )
        .into_response()
}
