//! Per-call debug trace system.
//!
//! Every API call gets a structured JSON trace file written to
//! `~/.config/zerogravity/traces/{YYYY-MM-DD}/{HH-MM-SS}_{cascade_short}.json`.
//!
//! Designed for LLM consumption: compact, structured, no raw bodies.

use serde::Serialize;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::Mutex;

/// Shared trace state for `AppState`.
#[derive(Clone)]
pub struct TraceCollector {
    enabled: bool,
    traces_dir: PathBuf,
}

impl TraceCollector {
    pub fn new(enabled: bool) -> Self {
        let traces_dir = crate::platform::Platform::detect().traces_dir;
        Self {
            enabled,
            traces_dir,
        }
    }

    /// Whether tracing is enabled.
    pub fn enabled(&self) -> bool {
        self.enabled
    }

    /// Start a new trace for an API call. Returns `None` if tracing is disabled.
    pub fn start(
        &self,
        cascade_id: &str,
        endpoint: &str,
        model: &str,
        stream: bool,
    ) -> Option<TraceHandle> {
        if !self.enabled {
            return None;
        }
        let now = chrono::Utc::now();
        Some(TraceHandle {
            inner: Arc::new(Mutex::new(TraceData {
                cascade_id: cascade_id.to_string(),
                endpoint: endpoint.to_string(),
                model: model.to_string(),
                stream,
                started_at: now.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(),
                finished_at: None,
                duration_ms: 0,
                outcome: "in_progress".to_string(),
                client_request: None,
                turns: Vec::new(),
                usage: None,
                errors: Vec::new(),
            })),
            traces_dir: self.traces_dir.clone(),
            started_at_chrono: now,
            started_instant: Instant::now(),
        })
    }

    /// Delete trace directories older than `max_age_days`.
    pub fn cleanup_old_traces(&self, max_age_days: u32) {
        if !self.enabled {
            return;
        }
        let Ok(entries) = std::fs::read_dir(&self.traces_dir) else {
            return;
        };
        let cutoff = chrono::Utc::now() - chrono::Duration::days(max_age_days as i64);
        let cutoff_str = cutoff.format("%Y-%m-%d").to_string();

        for entry in entries.flatten() {
            let name = entry.file_name().to_string_lossy().to_string();
            // Directory names are YYYY-MM-DD — lexicographic comparison works
            if name < cutoff_str {
                if let Err(e) = std::fs::remove_dir_all(entry.path()) {
                    tracing::warn!(dir = %name, error = %e, "trace: failed to remove old trace dir");
                } else {
                    tracing::info!(dir = %name, "trace: cleaned up old trace dir");
                }
            }
        }
    }
}

/// Handle to a single in-flight trace. Cloneable, thread-safe.
#[derive(Clone)]
pub struct TraceHandle {
    inner: Arc<Mutex<TraceData>>,
    traces_dir: PathBuf,
    started_at_chrono: chrono::DateTime<chrono::Utc>,
    started_instant: Instant,
}

impl std::fmt::Debug for TraceHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TraceHandle")
            .field("traces_dir", &self.traces_dir)
            .finish()
    }
}

impl TraceHandle {
    /// Record the client request summary.
    pub async fn set_client_request(&self, req: ClientRequestSummary) {
        self.inner.lock().await.client_request = Some(req);
    }

    /// Start a new turn (thinking/tool-call round). Returns the turn index.
    pub async fn start_turn(&self) -> usize {
        let mut data = self.inner.lock().await;
        let idx = data.turns.len();
        data.turns.push(TraceTurn {
            turn: idx,
            mitm_matched: false,
            gate_wait_ms: None,
            modify_summary: None,
            request_bytes: None,
            upstream_wait_ms: None,
            response: None,
            events_sent: Vec::new(),
            handler_action: None,
        });
        idx
    }

    /// Record MITM match for a turn.
    pub async fn record_mitm_match(&self, turn: usize, gate_wait_ms: u64) {
        let mut data = self.inner.lock().await;
        if let Some(t) = data.turns.get_mut(turn) {
            t.mitm_matched = true;
            t.gate_wait_ms = Some(gate_wait_ms);
        }
    }

    /// Record MITM modify summary for a turn.
    pub async fn record_modify(&self, turn: usize, summary: String, original: u64, modified: u64) {
        let mut data = self.inner.lock().await;
        if let Some(t) = data.turns.get_mut(turn) {
            t.modify_summary = Some(summary);
            t.request_bytes = Some(RequestBytes { original, modified });
        }
    }

    /// Record upstream wait time.
    pub async fn record_upstream_wait(&self, turn: usize, wait_ms: u64) {
        let mut data = self.inner.lock().await;
        if let Some(t) = data.turns.get_mut(turn) {
            t.upstream_wait_ms = Some(wait_ms);
        }
    }

    /// Record the response summary for a turn.
    pub async fn record_response(&self, turn: usize, resp: ResponseSummary) {
        let mut data = self.inner.lock().await;
        if let Some(t) = data.turns.get_mut(turn) {
            t.response = Some(resp);
        }
    }

    /// Record an event sent via channel.
    pub async fn record_event(&self, turn: usize, event_name: &str) {
        let mut data = self.inner.lock().await;
        if let Some(t) = data.turns.get_mut(turn) {
            t.events_sent.push(event_name.to_string());
        }
    }

    /// Record the handler action for a turn.
    pub async fn record_action(&self, turn: usize, action: &str) {
        let mut data = self.inner.lock().await;
        if let Some(t) = data.turns.get_mut(turn) {
            t.handler_action = Some(action.to_string());
        }
    }

    /// Record an error.
    pub async fn record_error(&self, error: String) {
        self.inner.lock().await.errors.push(error);
    }

    /// Record final usage.
    pub async fn set_usage(&self, usage: TrackedUsage) {
        self.inner.lock().await.usage = Some(usage);
    }

    /// Finalize the trace and write to disk as a per-call folder.
    ///
    /// Layout: `traces/{YYYY-MM-DD}/{HH-MM-SS}_{cascade_short}/`
    ///   - `summary.md`  — always written, rich LLM-readable overview
    ///   - `request.json` — client request details (always)
    ///   - `turns.json`   — per-turn MITM/gate/modify data (always)
    ///   - `response.json` — response summary + usage (if present)
    ///   - `events.json`  — channel events (if non-empty)
    ///   - `errors.json`  — errors (if any)
    pub async fn finish(&self, outcome: &str) {
        let mut data = self.inner.lock().await;
        let now = chrono::Utc::now();
        data.finished_at = Some(now.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string());
        data.duration_ms = self.started_instant.elapsed().as_millis() as u64;
        data.outcome = outcome.to_string();

        // Build folder path: traces/{YYYY-MM-DD}/{HH-MM-SS}_{cascade_short}/
        let date_str = self.started_at_chrono.format("%Y-%m-%d").to_string();
        let time_str = self.started_at_chrono.format("%H-%M-%S%.3f").to_string();
        let cascade_short = &data.cascade_id[..8.min(data.cascade_id.len())];
        let dir = self
            .traces_dir
            .join(&date_str)
            .join(format!("{}_{}", time_str, cascade_short));

        // Build all file contents while holding lock
        let summary = generate_summary(&data);
        let request_json = serde_json::to_string_pretty(&data.client_request).unwrap_or_default();
        let turns_json = serde_json::to_string_pretty(&data.turns).unwrap_or_default();

        let response_json =
            if data.usage.is_some() || data.turns.iter().any(|t| t.response.is_some()) {
                let resp = ResponseFile {
                    usage: data.usage.clone(),
                };
                Some(serde_json::to_string_pretty(&resp).unwrap_or_default())
            } else {
                None
            };

        let events_json = {
            let all_events: Vec<_> = data
                .turns
                .iter()
                .enumerate()
                .filter(|(_, t)| !t.events_sent.is_empty())
                .map(|(i, t)| serde_json::json!({ "turn": i, "events": t.events_sent }))
                .collect();
            if all_events.is_empty() {
                None
            } else {
                Some(serde_json::to_string_pretty(&all_events).unwrap_or_default())
            }
        };

        let errors_json = if data.errors.is_empty() {
            None
        } else {
            Some(serde_json::to_string_pretty(&data.errors).unwrap_or_default())
        };

        // Build meta.txt for grep
        let meta_txt = format!(
            "cascade={} endpoint={} model={} outcome={} duration={}ms stream={}",
            cascade_short, data.endpoint, data.model, data.outcome, data.duration_ms, data.stream
        );

        drop(data); // release lock before I/O

        tokio::spawn(async move {
            if let Err(e) = tokio::fs::create_dir_all(&dir).await {
                tracing::warn!(error = %e, "trace: failed to create dir");
                return;
            }
            // Always write summary + request + turns + meta
            let _ = tokio::fs::write(dir.join("summary.md"), summary).await;
            let _ = tokio::fs::write(dir.join("request.json"), request_json).await;
            let _ = tokio::fs::write(dir.join("turns.json"), turns_json).await;
            let _ = tokio::fs::write(dir.join("meta.txt"), meta_txt).await;
            // Conditionally write response, events, errors
            if let Some(j) = response_json {
                let _ = tokio::fs::write(dir.join("response.json"), j).await;
            }
            if let Some(j) = events_json {
                let _ = tokio::fs::write(dir.join("events.json"), j).await;
            }
            if let Some(j) = errors_json {
                let _ = tokio::fs::write(dir.join("errors.json"), j).await;
            }
            tracing::debug!(path = %dir.display(), "trace: folder written");
        });
    }
}

// ── Summary generation ─────────────────────────────────────────────────

#[derive(Serialize)]
struct ResponseFile {
    #[serde(skip_serializing_if = "Option::is_none")]
    usage: Option<TrackedUsage>,
}

/// Build a rich markdown summary from trace data.
fn generate_summary(data: &TraceData) -> String {
    let mut s = String::with_capacity(2048);
    let cascade_short = &data.cascade_id[..8.min(data.cascade_id.len())];

    // Header
    s.push_str(&format!(
        "# Trace: {} — {}\n\n",
        cascade_short, data.endpoint
    ));

    // Overview table
    s.push_str("| Field | Value |\n|-------|-------|\n");
    s.push_str(&format!("| Cascade ID | `{}` |\n", data.cascade_id));
    s.push_str(&format!("| Model | {} |\n", data.model));
    s.push_str(&format!("| Stream | {} |\n", data.stream));
    s.push_str(&format!("| Started | {} |\n", data.started_at));
    if let Some(ref fin) = data.finished_at {
        s.push_str(&format!("| Finished | {} |\n", fin));
    }
    s.push_str(&format!("| Duration | {}ms |\n", data.duration_ms));
    s.push_str(&format!("| Outcome | **{}** |\n", data.outcome));
    s.push('\n');

    // Client request
    s.push_str("## Client Request\n\n");
    if let Some(ref req) = data.client_request {
        s.push_str(&format!(
            "- **Messages:** {} (user text: {} chars)\n",
            req.message_count, req.user_text_len
        ));
        if !req.user_text_preview.is_empty() {
            s.push_str(&format!("- **Preview:** `{}`\n", req.user_text_preview));
        }
        s.push_str(&format!(
            "- **Tools:** {} | **Tool rounds:** {}\n",
            req.tool_count, req.tool_round_count
        ));
        if req.system_prompt {
            s.push_str("- **System prompt:** yes\n");
        }
        s.push_str(&format!(
            "- **Image:** {}\n",
            if req.has_image { "yes" } else { "no" }
        ));
    } else {
        s.push_str("(not recorded)\n");
    }
    s.push_str("\n→ Full details in [request.json](./request.json)\n\n");

    // Turns
    s.push_str(&format!("## Turns ({} total)\n\n", data.turns.len()));
    for turn in &data.turns {
        s.push_str(&format!("### Turn {}\n\n", turn.turn));

        // MITM match
        if turn.mitm_matched {
            s.push_str(&format!(
                "- **MITM matched:** ✓ (gate wait: {}ms)\n",
                turn.gate_wait_ms.unwrap_or(0)
            ));
        } else {
            s.push_str("- **MITM matched:** ✗\n");
        }

        // Modify
        if let Some(ref mod_sum) = turn.modify_summary {
            s.push_str(&format!("- **Modify:** {}", mod_sum));
            if let Some(ref bytes) = turn.request_bytes {
                s.push_str(&format!(" ({}B → {}B)", bytes.original, bytes.modified));
            }
            s.push('\n');
        }

        // Upstream wait
        if let Some(wait) = turn.upstream_wait_ms {
            s.push_str(&format!("- **Upstream wait:** {}ms\n", wait));
        }

        // Response
        if let Some(ref resp) = turn.response {
            s.push_str(&format!(
                "- **Response:** {} chars text, {} chars thinking",
                resp.text_len, resp.thinking_len
            ));
            if let Some(ref fr) = resp.finish_reason {
                s.push_str(&format!(", finish_reason={}", fr));
            }
            if !resp.function_calls.is_empty() {
                let names: Vec<&str> = resp
                    .function_calls
                    .iter()
                    .map(|f| f.name.as_str())
                    .collect();
                s.push_str(&format!(", tool_calls=[{}]", names.join(", ")));
            }
            if resp.grounding {
                s.push_str(", grounding=yes");
            }
            s.push('\n');
            if !resp.text_preview.is_empty() {
                s.push_str(&format!("- **Output preview:** `{}`\n", resp.text_preview));
            }
        }

        // Events
        if !turn.events_sent.is_empty() {
            s.push_str(&format!(
                "- **Events:** {} sent ({})\n",
                turn.events_sent.len(),
                turn.events_sent.join(", ")
            ));
        }

        // Handler action
        if let Some(ref action) = turn.handler_action {
            s.push_str(&format!("- **Action:** {}\n", action));
        }

        s.push('\n');
    }

    if data.turns.iter().any(|t| t.response.is_some()) {
        s.push_str("→ Full turn details in [turns.json](./turns.json)\n\n");
    }

    // Usage
    if let Some(ref u) = data.usage {
        s.push_str("## Usage\n\n");
        s.push_str("| Metric | Tokens |\n|--------|--------|\n");
        s.push_str(&format!("| Input | {} |\n", u.input_tokens));
        s.push_str(&format!("| Output | {} |\n", u.output_tokens));
        if u.thinking_tokens > 0 {
            s.push_str(&format!("| Thinking | {} |\n", u.thinking_tokens));
        }
        if u.cache_read > 0 {
            s.push_str(&format!("| Cache read | {} |\n", u.cache_read));
        }
        s.push_str("\n→ Full details in [response.json](./response.json)\n\n");
    }

    // Errors
    if !data.errors.is_empty() {
        s.push_str("## Errors\n\n");
        for err in &data.errors {
            s.push_str(&format!("- ❌ {}\n", err));
        }
        s.push_str("\n→ Full details in [errors.json](./errors.json)\n\n");
    }

    // Files index
    s.push_str("## Files\n\n");
    s.push_str("| File | Contains |\n|------|----------|\n");
    s.push_str("| [request.json](./request.json) | Client request summary |\n");
    s.push_str("| [turns.json](./turns.json) | Per-turn MITM/gate/modify/response data |\n");
    if data.usage.is_some() || data.turns.iter().any(|t| t.response.is_some()) {
        s.push_str("| [response.json](./response.json) | Response summaries + token usage |\n");
    }
    if data.turns.iter().any(|t| !t.events_sent.is_empty()) {
        s.push_str("| [events.json](./events.json) | Channel events per turn |\n");
    }
    if !data.errors.is_empty() {
        s.push_str("| [errors.json](./errors.json) | Error messages |\n");
    }

    s
}

// ── Serializable data structures ───────────────────────────────────────

#[derive(Serialize)]
struct TraceData {
    cascade_id: String,
    endpoint: String,
    model: String,
    stream: bool,
    started_at: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    finished_at: Option<String>,
    duration_ms: u64,
    outcome: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    client_request: Option<ClientRequestSummary>,
    turns: Vec<TraceTurn>,
    #[serde(skip_serializing_if = "Option::is_none")]
    usage: Option<TrackedUsage>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    errors: Vec<String>,
}

#[derive(Serialize, Clone)]
pub struct ClientRequestSummary {
    pub message_count: usize,
    pub tool_count: usize,
    pub tool_round_count: usize,
    pub user_text_len: usize,
    pub user_text_preview: String,
    pub system_prompt: bool,
    pub has_image: bool,
}

#[derive(Serialize)]
struct TraceTurn {
    turn: usize,
    mitm_matched: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    gate_wait_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    modify_summary: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    request_bytes: Option<RequestBytes>,
    #[serde(skip_serializing_if = "Option::is_none")]
    upstream_wait_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    response: Option<ResponseSummary>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    events_sent: Vec<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    handler_action: Option<String>,
}

#[derive(Serialize)]
struct RequestBytes {
    original: u64,
    modified: u64,
}

#[derive(Serialize, Clone)]
pub struct ResponseSummary {
    pub text_len: usize,
    pub thinking_len: usize,
    #[serde(skip_serializing_if = "String::is_empty")]
    pub text_preview: String,
    pub finish_reason: Option<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub function_calls: Vec<FunctionCallSummary>,
    pub grounding: bool,
}

#[derive(Serialize, Clone)]
pub struct FunctionCallSummary {
    pub name: String,
    pub args_preview: String,
}

#[derive(Serialize, Clone)]
pub struct TrackedUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub thinking_tokens: u64,
    pub cache_read: u64,
}
