//! Shared polling engine and step extraction helpers.

use rand::Rng;
use tracing::{debug, info, warn};

use super::AppState;

/// Real token usage reported by the language server.
#[derive(Debug, Clone, Default)]
#[allow(dead_code)]
pub(crate) struct ModelUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub api_provider: String,
    pub model: String,
}

/// Result of polling — text + optional real usage data + thinking data.
pub(crate) struct PollResult {
    pub text: String,
    pub usage: Option<ModelUsage>,
    /// Opaque Anthropic thinking verification signature from PLANNER_RESPONSE.
    /// Required for multi-turn thinking model chaining.
    pub thinking_signature: Option<String>,
    /// The model's internal reasoning/thinking content.
    /// Available for both Opus (Anthropic) and Gemini models.
    pub thinking: Option<String>,
    /// Time the model spent thinking, as reported by the LS (e.g. "0.041999832s").
    #[allow(dead_code)]
    pub thinking_duration: Option<String>,
    /// Upstream error from Google's API, captured by MITM.
    pub upstream_error: Option<crate::mitm::store::UpstreamError>,
}

/// Extract the response text from steps — scans in REVERSE to find the latest response.
pub(crate) fn extract_response_text(steps: &[serde_json::Value]) -> String {
    for step in steps.iter().rev() {
        let step_type = step["type"].as_str().unwrap_or("");

        if step_type.contains("PLANNER_RESPONSE") {
            let resp = &step["plannerResponse"];
            let text = resp["rawResponse"]
                .as_str()
                .or_else(|| resp["response"].as_str())
                .unwrap_or("");
            if !text.is_empty() {
                return text.to_string();
            }
        }

        if step_type.contains("AI_RESPONSE") || step_type.contains("MODEL_RESPONSE") {
            if let Some(text) = step["response"]
                .as_str()
                .or_else(|| step["rawResponse"].as_str())
                .or_else(|| step["text"].as_str())
            {
                if !text.is_empty() {
                    return text.to_string();
                }
            }
        }
    }
    String::new()
}

/// Extract real token usage from the LS's modelUsage field.
/// The LS reports this in CHECKPOINT steps and sometimes in retryInfos.
/// Scans in reverse to find the latest usage data.
pub(crate) fn extract_model_usage(steps: &[serde_json::Value]) -> Option<ModelUsage> {
    for step in steps.iter().rev() {
        if let Some(usage) = step.get("metadata").and_then(|m| m.get("modelUsage")) {
            let input = usage["inputTokens"]
                .as_str()
                .and_then(|s| s.parse::<u64>().ok())
                .or_else(|| usage["inputTokens"].as_u64())
                .unwrap_or(0);
            let output = usage["outputTokens"]
                .as_str()
                .and_then(|s| s.parse::<u64>().ok())
                .or_else(|| usage["outputTokens"].as_u64())
                .unwrap_or(0);

            if input > 0 || output > 0 {
                return Some(ModelUsage {
                    input_tokens: input,
                    output_tokens: output,
                    api_provider: usage["apiProvider"].as_str().unwrap_or("").to_string(),
                    model: usage["model"].as_str().unwrap_or("").to_string(),
                });
            }
        }
    }
    None
}

/// Extract the thinking signature from PLANNER_RESPONSE steps.
/// This is an opaque Base64 blob used by Anthropic for extended thinking
/// verification. Needed to chain multi-turn conversations with thinking models.
pub(crate) fn extract_thinking_signature(steps: &[serde_json::Value]) -> Option<String> {
    for step in steps.iter().rev() {
        let step_type = step["type"].as_str().unwrap_or("");
        if step_type.contains("PLANNER_RESPONSE") {
            if let Some(sig) = step["plannerResponse"]["thinkingSignature"].as_str() {
                if !sig.is_empty() {
                    return Some(sig.to_string());
                }
            }
        }
    }
    None
}

/// Extract the model's thinking/reasoning content from PLANNER_RESPONSE steps.
/// This is the internal monologue produced during extended thinking.
/// Available for ALL models (Opus, Gemini Flash, Gemini Pro).
pub(crate) fn extract_thinking_content(steps: &[serde_json::Value]) -> Option<String> {
    for step in steps.iter().rev() {
        let step_type = step["type"].as_str().unwrap_or("");
        if step_type.contains("PLANNER_RESPONSE") {
            if let Some(thinking) = step["plannerResponse"]["thinking"].as_str() {
                if !thinking.is_empty() {
                    return Some(thinking.to_string());
                }
            }
        }
    }
    None
}

/// Extract thinking duration from PLANNER_RESPONSE steps.
/// Returns the raw duration string as reported by the LS (e.g. "0.041999832s").
pub(crate) fn extract_thinking_duration(steps: &[serde_json::Value]) -> Option<String> {
    for step in steps.iter().rev() {
        let step_type = step["type"].as_str().unwrap_or("");
        if step_type.contains("PLANNER_RESPONSE") {
            if let Some(dur) = step["plannerResponse"]["thinkingDuration"].as_str() {
                if !dur.is_empty() {
                    return Some(dur.to_string());
                }
            }
        }
    }
    None
}

/// Check if the cascade has truly finished — the last PLANNER_RESPONSE must be DONE
/// AND the very last step must be a terminal type (CHECKPOINT or PLANNER_RESPONSE with DONE).
/// This prevents false positives during agentic tool-call loops where intermediate
/// PLANNER_RESPONSE steps show DONE but the cascade keeps going.
pub(crate) fn is_response_done(steps: &[serde_json::Value]) -> bool {
    if steps.is_empty() {
        return false;
    }

    let last = &steps[steps.len() - 1];
    let last_type = last["type"].as_str().unwrap_or("");
    let last_status = last["status"].as_str().unwrap_or("");

    // CHECKPOINT at the end = cascade is definitely done
    if last_type.contains("CHECKPOINT") {
        return true;
    }

    // Last step is a PLANNER_RESPONSE with DONE = final answer (no more tool calls coming)
    if (last_type.contains("PLANNER_RESPONSE")
        || last_type.contains("AI_RESPONSE")
        || last_type.contains("MODEL_RESPONSE"))
        && last_status.contains("DONE")
    {
        return true;
    }

    false
}

/// Poll the backend until we get a response or timeout.
pub(crate) async fn poll_for_response(
    state: &AppState,
    cascade_id: &str,
    timeout: u64,
) -> PollResult {
    let start = std::time::Instant::now();
    let short_id = &cascade_id[..8.min(cascade_id.len())];
    info!("Polling for response on cascade {short_id} (timeout={timeout}s)");

    let mut last_step_count: usize = 0;

    while start.elapsed().as_secs() < timeout {
        // Check for upstream errors from MITM (Google API errors)
        if let Some(err) = state.mitm_store.take_upstream_error().await {
            warn!(
                "Upstream error on cascade {short_id}: HTTP {} — {}",
                err.status,
                err.message.as_deref().unwrap_or("unknown")
            );
            return PollResult {
                text: String::new(),
                usage: None,
                thinking_signature: None,
                thinking: None,
                thinking_duration: None,
                upstream_error: Some(err),
            };
        }

        if let Ok((status, data)) = state.backend.get_steps(cascade_id).await {
            if status == 200 {
                if let Some(steps) = data["steps"].as_array() {
                    let step_count = steps.len();

                    // Only log when step count changes (denoised)
                    if step_count != last_step_count {
                        // Compact type summary: count unique types
                        let mut type_counts: std::collections::BTreeMap<&str, usize> =
                            std::collections::BTreeMap::new();
                        for s in steps.iter() {
                            let t = s["type"]
                                .as_str()
                                .unwrap_or("?")
                                .strip_prefix("CORTEX_STEP_TYPE_")
                                .unwrap_or("?");
                            *type_counts.entry(t).or_insert(0) += 1;
                        }
                        let summary: Vec<String> = type_counts
                            .iter()
                            .map(|(t, c)| {
                                if *c > 1 {
                                    format!("{t}×{c}")
                                } else {
                                    t.to_string()
                                }
                            })
                            .collect();
                        debug!(
                            "Poll {short_id}: {step_count} steps [{}]",
                            summary.join(", ")
                        );
                        last_step_count = step_count;
                    }

                    // Check if the cascade is truly done
                    if is_response_done(steps) {
                        let text = extract_response_text(steps);
                        if !text.is_empty() {
                            let usage = extract_model_usage(steps);
                            let thinking_signature = extract_thinking_signature(steps);
                            let thinking = extract_thinking_content(steps);
                            let thinking_duration = extract_thinking_duration(steps);
                            let elapsed = start.elapsed().as_secs_f32();
                            if let Some(ref u) = usage {
                                info!(
                                    "Response done ({short_id}), {:.1}s, {} chars, tokens: {}in/{}out ({}){}{}",
                                    elapsed, text.len(), u.input_tokens, u.output_tokens, u.model,
                                    thinking.as_ref().map_or(String::new(), |t| format!(", thinking: {} chars", t.len())),
                                    if thinking_signature.is_some() { ", has sig" } else { "" }
                                );
                            } else {
                                info!(
                                    "Response done ({short_id}), {:.1}s, {} chars (no usage){}{}",
                                    elapsed,
                                    text.len(),
                                    thinking.as_ref().map_or(String::new(), |t| format!(
                                        ", thinking: {} chars",
                                        t.len()
                                    )),
                                    if thinking_signature.is_some() {
                                        ", has sig"
                                    } else {
                                        ""
                                    }
                                );
                            }
                            return PollResult {
                                text,
                                usage,
                                thinking_signature,
                                thinking,
                                thinking_duration,
                                upstream_error: None,
                            };
                        }
                    }

                    // Fallback: check trajectory IDLE status (catches edge cases)
                    // Only check every 5th poll to reduce network calls
                    if step_count > 4 && step_count % 5 == 0 {
                        if let Ok((ts, td)) = state.backend.get_trajectory(cascade_id).await {
                            if ts == 200 {
                                let run_status = td["status"].as_str().unwrap_or("");
                                if run_status.contains("IDLE") {
                                    let text = extract_response_text(steps);
                                    if !text.is_empty() {
                                        let usage = extract_model_usage(steps);
                                        let thinking_signature = extract_thinking_signature(steps);
                                        let thinking = extract_thinking_content(steps);
                                        let thinking_duration = extract_thinking_duration(steps);
                                        let elapsed = start.elapsed().as_secs_f32();
                                        info!(
                                            "Trajectory IDLE ({short_id}), {:.1}s, {} chars",
                                            elapsed,
                                            text.len()
                                        );
                                        return PollResult {
                                            text,
                                            usage,
                                            thinking_signature,
                                            thinking,
                                            thinking_duration,
                                            upstream_error: None,
                                        };
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

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

    warn!("Timeout after {timeout}s on cascade {short_id}");
    PollResult {
        text: String::new(),
        usage: None,
        thinking_signature: None,
        thinking: None,
        thinking_duration: None,
        upstream_error: Some(crate::mitm::store::UpstreamError {
            status: 504,
            body: String::new(),
            message: Some(format!(
                "Timeout: no response from Google API after {timeout}s"
            )),
            error_status: Some("DEADLINE_EXCEEDED".to_string()),
        }),
    }
}
