//! Pure search endpoint (/v1/search) — triggers Google Search via the LS.
//!
//! Sends a minimal prompt to the LS with Google Search grounding enabled,
//! captures the grounding metadata from the response, and returns the
//! search results without the model's generated text.
//!
//! The LS triggers the actual search request to Google's servers;
//! we capture the results via MITM, never calling Google directly.

use axum::{
    extract::State,
    http::StatusCode,
    response::{IntoResponse, Json},
};
use std::sync::Arc;
use tracing::{info, warn};

use super::models::{lookup_model, MODELS};
use super::polling::poll_for_response;
use super::util::{err_response, upstream_err_response};
use super::AppState;

/// Search request body.
#[derive(serde::Deserialize)]
pub(crate) struct SearchRequest {
    /// Search query.
    pub query: String,
    /// Model to use for grounding. Defaults to gemini-3-flash (cheapest).
    #[serde(default = "default_search_model")]
    pub model: String,
    /// Timeout in seconds.
    #[serde(default = "default_search_timeout")]
    pub timeout: u64,
    /// Conversation/session ID for context reuse.
    #[serde(default)]
    #[allow(dead_code)]
    pub conversation: Option<String>,
    /// Max output tokens — keep low since we only want grounding metadata.
    #[serde(default = "default_search_max_tokens")]
    pub max_output_tokens: u64,
}

fn default_search_model() -> String {
    "gemini-3-flash".to_string()
}

fn default_search_timeout() -> u64 {
    30
}

fn default_search_max_tokens() -> u64 {
    256
}

/// GET /v1/search?q=... — quick search via query param.
pub(crate) async fn handle_search_get(
    State(state): State<Arc<AppState>>,
    axum::extract::Query(params): axum::extract::Query<SearchQueryParams>,
) -> axum::response::Response {
    info!("GET /v1/search q={}", params.q);

    let body = SearchRequest {
        query: params.q,
        model: params.model.unwrap_or_else(default_search_model),
        timeout: params.timeout.unwrap_or(default_search_timeout()),
        conversation: None,
        max_output_tokens: params.max_tokens.unwrap_or(default_search_max_tokens()),
    };

    do_search(state, body).await
}

#[derive(serde::Deserialize)]
pub(crate) struct SearchQueryParams {
    pub q: String,
    #[serde(default)]
    pub model: Option<String>,
    #[serde(default)]
    pub timeout: Option<u64>,
    #[serde(default)]
    pub max_tokens: Option<u64>,
}

/// POST /v1/search — full search request.
pub(crate) async fn handle_search_post(
    State(state): State<Arc<AppState>>,
    Json(body): Json<SearchRequest>,
) -> axum::response::Response {
    info!("POST /v1/search q={}", body.query);
    do_search(state, body).await
}

async fn do_search(state: Arc<AppState>, body: SearchRequest) -> axum::response::Response {
    let model = match lookup_model(&body.model) {
        Some(m) => m,
        None => {
            let names: Vec<&str> = MODELS.iter().map(|m| m.name).collect();
            return err_response(
                StatusCode::BAD_REQUEST,
                format!("Unknown model: {}. Available: {names:?}", body.model),
                "invalid_request_error",
            );
        }
    };

    let token = state.backend.oauth_token().await;
    if token.is_empty() {
        return err_response(
            StatusCode::UNAUTHORIZED,
            "No OAuth token.".into(),
            "authentication_error",
        );
    }

    // Build generation params with Google Search grounding enabled
    use crate::mitm::store::GenerationParams;
    let gp = GenerationParams {
        max_output_tokens: Some(body.max_output_tokens),
        google_search: true,
        ..Default::default()
    };

    // Create a prompt that encourages the model to ground its response
    let search_prompt = format!(
        "Search the web for the following query and provide a brief summary of the results:\n\n{}",
        body.query
    );

    // 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::INTERNAL_SERVER_ERROR,
                format!("Failed to create cascade: {e}"),
                "server_error",
            );
        }
    };

    // Start debug trace
    let trace = state
        .trace
        .start(&cascade_id, "POST /v1/search", model.name, false);
    if let Some(ref t) = trace {
        t.set_client_request(crate::trace::ClientRequestSummary {
            message_count: 1,
            tool_count: 0,
            tool_round_count: 0,
            user_text_len: body.query.len(),
            user_text_preview: body.query.chars().take(200).collect(),
            system_prompt: false,
            has_image: false,
        })
        .await;
        t.start_turn().await;
    }

    let mitm_gate = std::sync::Arc::new(tokio::sync::Notify::new());
    let mitm_gate_clone = mitm_gate.clone();
    let (mitm_tx, mut mitm_rx) = tokio::sync::mpsc::channel(64);
    state
        .mitm_store
        .register_request(crate::mitm::store::RequestContext {
            cascade_id: cascade_id.clone(),
            pending_user_text: search_prompt.clone(),
            event_channel: mitm_tx,
            generation_params: Some(gp.clone()),
            pending_image: None,
            tools: None,
            tool_config: None,
            pending_tool_results: Vec::new(),
            tool_rounds: Vec::new(),
            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 dot to LS — real search prompt injected by MITM proxy
    if let Err(e) = state
        .backend
        .send_message(
            &cascade_id,
            &format!(".<cid:{}>", cascade_id),
            model.model_enum,
        )
        .await
    {
        state.mitm_store.remove_request(&cascade_id).await;
        return err_response(
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to send search message: {e}"),
            "server_error",
        );
    }

    // ── Strict timeout cascade ───────────────────────────────────────────────
    // 5s gate → MITM didn't match → 502
    let gate_matched =
        tokio::time::timeout(std::time::Duration::from_secs(5), mitm_gate.notified()).await;

    if gate_matched.is_err() {
        if state.mitm_enabled {
            state.mitm_store.remove_request(&cascade_id).await;
            return err_response(
                StatusCode::BAD_GATEWAY,
                "MITM proxy did not match request within 5s".to_string(),
                "mitm_timeout",
            );
        }
        // --no-mitm fallback: use polling
        tracing::warn!(cascade = %cascade_id, "MITM gate timeout (--no-mitm mode, falling back to polling)");
    }

    // ── Channel-based response path (primary) ────────────────────────────────
    if state.mitm_enabled {
        let timeout = body.timeout;
        let mut response_text = String::new();
        let mut last_usage: Option<crate::mitm::store::ApiUsage> = None;
        let mut retries = 0u32;
        const MAX_RETRIES: u32 = 3;

        while let Some(event) =
            tokio::time::timeout(std::time::Duration::from_secs(timeout), mitm_rx.recv())
                .await
                .ok()
                .flatten()
        {
            use crate::mitm::store::MitmEvent;
            match event {
                MitmEvent::TextDelta(t) => {
                    response_text.push_str(&t);
                }
                MitmEvent::ThinkingDelta(_) => {} // search doesn't use thinking
                MitmEvent::Usage(u) => {
                    last_usage = Some(u);
                }
                MitmEvent::Grounding(_) => {} // stored by proxy directly
                MitmEvent::FunctionCall(_) => {} // not expected for search
                MitmEvent::ResponseComplete => {
                    // Check if we got actual content — if not, this was a
                    // thinking-only intermediate response. The LS will make
                    // a follow-up request; re-register context and keep waiting.
                    let grounding_peek = state.mitm_store.peek_grounding().await;
                    if response_text.is_empty() && grounding_peek.is_none() {
                        retries += 1;
                        if retries >= MAX_RETRIES {
                            tracing::warn!(cascade = %cascade_id, retries, "Search: max retries reached with no content — giving up");
                            break;
                        }
                        let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
                        let new_gate = std::sync::Arc::new(tokio::sync::Notify::new());
                        state
                            .mitm_store
                            .register_request(crate::mitm::store::RequestContext {
                                cascade_id: cascade_id.clone(),
                                pending_user_text: search_prompt.clone(),
                                event_channel: new_tx,
                                generation_params: Some(gp.clone()),
                                pending_image: None,
                                tools: None,
                                tool_config: None,
                                pending_tool_results: Vec::new(),
                                tool_rounds: Vec::new(),
                                last_function_calls: Vec::new(),
                                call_id_to_name: std::collections::HashMap::new(),
                                created_at: std::time::Instant::now(),
                                gate: new_gate,
                                trace_handle: trace.clone(),
                                trace_turn: 0,
                            })
                            .await;
                        mitm_rx = new_rx;
                        tracing::debug!(
                            cascade = %cascade_id, retries,
                            "Search: empty response — re-registered context for follow-up"
                        );
                        continue;
                    }
                    break;
                }
                MitmEvent::UpstreamError(err) => {
                    if let Some(ref t) = trace {
                        t.record_error(format!(
                            "Upstream: {}",
                            super::util::upstream_error_message(&err)
                        ))
                        .await;
                        t.finish("upstream_error").await;
                    }
                    state.mitm_store.remove_request(&cascade_id).await;
                    return upstream_err_response(&err);
                }
            }
        }

        // Extract grounding metadata (stored by dispatch_stream_events)
        let grounding = state.mitm_store.take_grounding().await;
        state.mitm_store.remove_request(&cascade_id).await;

        if response_text.is_empty() && grounding.is_none() {
            if let Some(ref t) = trace {
                t.record_error(format!(
                    "Timeout: no search response after {timeout}s (retries: {retries})"
                ))
                .await;
                t.finish("timeout").await;
            }
            return err_response(
                StatusCode::GATEWAY_TIMEOUT,
                format!("Timeout: no search response after {timeout}s"),
                "upstream_error",
            );
        }

        return {
            // Finalize trace for channel-based path
            if let Some(ref t) = trace {
                t.record_response(
                    0,
                    crate::trace::ResponseSummary {
                        text_len: response_text.len(),
                        thinking_len: 0,
                        text_preview: response_text.chars().take(200).collect(),
                        finish_reason: Some("stop".to_string()),
                        function_calls: Vec::new(),
                        grounding: grounding.is_some(),
                    },
                )
                .await;
                if let Some((it, ot)) = last_usage
                    .as_ref()
                    .map(|u| (u.input_tokens, u.output_tokens))
                {
                    t.set_usage(crate::trace::TrackedUsage {
                        input_tokens: it,
                        output_tokens: ot,
                        thinking_tokens: 0,
                        cache_read: 0,
                    })
                    .await;
                }
                t.finish("completed").await;
            }
            build_search_response(
                &body.query,
                model.name,
                response_text,
                grounding,
                last_usage.map(|u| (u.input_tokens, u.output_tokens)),
            )
        };
    }

    // ── Fallback: polling path (--no-mitm only) ──────────────────────────────
    let poll_result = poll_for_response(&state, &cascade_id, body.timeout).await;
    if let Some(ref err) = poll_result.upstream_error {
        return upstream_err_response(err);
    }

    let grounding = state.mitm_store.take_grounding().await;

    let response_text = if !poll_result.text.is_empty() {
        poll_result.text.clone()
    } else {
        state
            .mitm_store
            .take_response_text()
            .await
            .unwrap_or_default()
    };

    state.mitm_store.remove_request(&cascade_id).await;
    state.mitm_store.clear_response_async().await;

    // Finalize trace for polling path
    if let Some(ref t) = trace {
        t.record_response(
            0,
            crate::trace::ResponseSummary {
                text_len: response_text.len(),
                thinking_len: 0,
                text_preview: response_text.chars().take(200).collect(),
                finish_reason: Some("stop".to_string()),
                function_calls: Vec::new(),
                grounding: grounding.is_some(),
            },
        )
        .await;
        t.finish("completed").await;
    }

    build_search_response(
        &body.query,
        model.name,
        response_text,
        grounding,
        poll_result.usage.map(|u| (u.input_tokens, u.output_tokens)),
    )
}

fn build_search_response(
    query: &str,
    model_name: &str,
    response_text: String,
    grounding: Option<serde_json::Value>,
    usage: Option<(u64, u64)>,
) -> axum::response::Response {
    use axum::Json;

    let mut response = serde_json::json!({
        "object": "search_result",
        "query": query,
        "model": model_name,
        "summary": response_text,
    });

    // Include grounding metadata if available
    if let Some(ref gm) = grounding {
        // Extract structured search results
        let mut search_results = Vec::new();

        // groundingChunks → individual web results
        if let Some(chunks) = gm.get("groundingChunks").and_then(|v| v.as_array()) {
            for chunk in chunks {
                if let Some(web) = chunk.get("web") {
                    search_results.push(serde_json::json!({
                        "title": web.get("title").and_then(|v| v.as_str()).unwrap_or(""),
                        "url": web.get("uri").and_then(|v| v.as_str()).unwrap_or(""),
                    }));
                }
            }
        }

        // groundingSupports → citations with source references
        let mut citations = Vec::new();
        if let Some(supports) = gm.get("groundingSupports").and_then(|v| v.as_array()) {
            for support in supports {
                let text = support
                    .get("segment")
                    .and_then(|s| s.get("text"))
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                let indices: Vec<u64> = support
                    .get("groundingChunkIndices")
                    .and_then(|v| v.as_array())
                    .map(|arr| arr.iter().filter_map(|i| i.as_u64()).collect())
                    .unwrap_or_default();
                let scores: Vec<f64> = support
                    .get("confidenceScores")
                    .and_then(|v| v.as_array())
                    .map(|arr| arr.iter().filter_map(|s| s.as_f64()).collect())
                    .unwrap_or_default();

                citations.push(serde_json::json!({
                    "text": text,
                    "source_indices": indices,
                    "confidence_scores": scores,
                }));
            }
        }

        // searchEntryPoint → rendered search widget HTML
        let search_url = gm
            .get("searchEntryPoint")
            .and_then(|sep| sep.get("renderedContent"))
            .and_then(|v| v.as_str());

        // webSearchQueries → the actual queries Google used
        let queries = gm
            .get("webSearchQueries")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|q| q.as_str().map(|s| s.to_string()))
                    .collect::<Vec<_>>()
            });

        response["results"] = serde_json::json!(search_results);
        response["citations"] = serde_json::json!(citations);
        if let Some(qs) = queries {
            response["search_queries"] = serde_json::json!(qs);
        }
        if let Some(url) = search_url {
            response["search_widget_html"] = serde_json::json!(url);
        }

        // Include raw grounding metadata for advanced consumers
        response["grounding_metadata"] = gm.clone();
    } else {
        response["results"] = serde_json::json!([]);
        response["citations"] = serde_json::json!([]);
        warn!("Search completed but no grounding metadata captured — model may not have grounded");
    }

    // Include usage if available
    if let Some((input, output)) = usage {
        response["usage"] = serde_json::json!({
            "input_tokens": input,
            "output_tokens": output,
            "total_tokens": input + output,
        });
    }

    Json(response).into_response()
}
