use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};

pub const TRACE_SCHEMA_VERSION: u32 = 1;
pub const TRACE_FILES: &[&str] = &[
    "runs.jsonl",
    "events.jsonl",
    "subagents.jsonl",
    "task-gates.jsonl",
];

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct TraceEvent {
    pub schema_version: u32,
    pub event_id: String,
    pub ts: String,
    pub run_id: String,
    pub parent_run_id: Option<String>,
    pub orchestration: Option<String>,
    pub stage: Option<String>,
    pub task_id: Option<String>,
    pub agent: Option<String>,
    pub provider: Option<String>,
    pub kind: String,
    pub status: Option<String>,
    pub tool: Option<String>,
    pub target: Option<String>,
    pub duration_ms: Option<u128>,
    pub usage: Option<Value>,
    pub context: Option<Value>,
    pub evidence: Option<Value>,
    pub error: Option<String>,
    pub metadata: Value,
}

#[derive(Clone, Debug, Default, Serialize)]
pub struct TraceSummary {
    pub total_events: usize,
    pub root_count: usize,
    pub roots: Vec<String>,
    pub by_kind: BTreeMap<String, usize>,
    pub by_status: BTreeMap<String, usize>,
    pub by_stage: BTreeMap<String, usize>,
    pub latest_ts: Option<String>,
}

#[derive(Clone, Debug, Serialize)]
pub struct TraceTree {
    pub run_id: String,
    pub events: Vec<TraceEvent>,
    pub children: Vec<TraceTree>,
}

#[derive(Clone, Debug, Serialize)]
pub struct TraceDoctorReport {
    pub status: &'static str,
    pub files: BTreeMap<String, usize>,
    pub warnings: Vec<String>,
}

impl TraceEvent {
    pub fn from_payload(payload: Value, source_file: &str, source_line: usize) -> Self {
        let payload = redact_value(payload);
        let object = payload.as_object().cloned();
        let map = object.as_ref();
        let ts = string_field(map, "ts")
            .or_else(|| string_field(map, "timestamp"))
            .unwrap_or_else(|| "unknown".to_string());
        let kind = string_field(map, "kind")
            .or_else(|| string_field(map, "event"))
            .or_else(|| string_field(map, "hook_event_name"))
            .unwrap_or_else(|| source_file.trim_end_matches(".jsonl").to_string());
        let status = string_field(map, "status").or_else(|| string_field(map, "decision"));
        let run_id = string_field(map, "run_id")
            .or_else(|| string_field(map, "session_id"))
            .or_else(|| string_field(map, "task_id").map(|task| format!("task:{task}")))
            .unwrap_or_else(|| legacy_run_id(source_file, source_line, &payload));
        let event_id = string_field(map, "event_id")
            .unwrap_or_else(|| event_id(source_file, source_line, &payload));
        let metadata = match object.clone() {
            Some(mut object) => {
                for key in [
                    "schema_version",
                    "event_id",
                    "ts",
                    "timestamp",
                    "run_id",
                    "parent_run_id",
                    "orchestration",
                    "stage",
                    "task_id",
                    "agent",
                    "agent_type",
                    "subagent_type",
                    "provider",
                    "kind",
                    "event",
                    "hook_event_name",
                    "status",
                    "decision",
                    "tool",
                    "tool_name",
                    "target",
                    "duration_ms",
                    "usage",
                    "context",
                    "evidence",
                    "error",
                ] {
                    object.remove(key);
                }
                Value::Object(object)
            }
            None => json!({ "raw": payload }),
        };
        let error = string_field(map, "error").or_else(|| {
            status
                .as_deref()
                .filter(|status| {
                    matches!(
                        *status,
                        "blocked" | "failed" | "failure" | "error" | "rejected"
                    )
                })
                .and_then(|_| string_field(map, "reason"))
        });
        Self {
            schema_version: TRACE_SCHEMA_VERSION,
            event_id,
            ts,
            run_id,
            parent_run_id: string_field(map, "parent_run_id"),
            orchestration: string_field(map, "orchestration"),
            stage: string_field(map, "stage"),
            task_id: string_field(map, "task_id"),
            agent: string_field(map, "agent")
                .or_else(|| string_field(map, "agent_type"))
                .or_else(|| string_field(map, "subagent_type")),
            provider: string_field(map, "provider"),
            kind,
            status,
            tool: string_field(map, "tool").or_else(|| string_field(map, "tool_name")),
            target: string_field(map, "target"),
            duration_ms: u128_field(map, "duration_ms"),
            usage: map.and_then(|object| object.get("usage").cloned()),
            context: map.and_then(|object| object.get("context").cloned()),
            evidence: map.and_then(|object| object.get("evidence").cloned()),
            error,
            metadata,
        }
    }

    pub fn to_value(&self) -> Value {
        serde_json::to_value(self).unwrap_or_else(|_| json!({}))
    }
}

pub fn make_event_id(seed: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(seed.as_bytes());
    let digest = hasher.finalize();
    format!("evt-{}", hex_prefix(&digest, 16))
}

pub fn make_run_id(seed: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(seed.as_bytes());
    let digest = hasher.finalize();
    format!("run-{}", hex_prefix(&digest, 16))
}

pub fn record_event(project_dir: &Path, file: &str, payload: Value) -> Result<TraceEvent> {
    let line = next_line(project_dir, file);
    let event = TraceEvent::from_payload(payload, file, line);
    append_event(project_dir, file, &event)?;
    if file != "runs.jsonl" {
        append_event(project_dir, "runs.jsonl", &event)?;
    }
    Ok(event)
}

pub fn append_event(project_dir: &Path, file: &str, event: &TraceEvent) -> Result<()> {
    let log_dir = project_dir.join(".sdd");
    fs::create_dir_all(&log_dir)?;
    let mut fh = fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(log_dir.join(file))?;
    writeln!(fh, "{}", serde_json::to_string(&event.to_value())?)?;
    Ok(())
}

pub fn read_events(project_dir: &Path) -> Result<Vec<TraceEvent>> {
    let mut events = Vec::new();
    let mut seen = BTreeSet::new();
    for file in TRACE_FILES {
        for event in read_events_from_file(project_dir, file)? {
            if seen.insert(event.event_id.clone()) {
                events.push(event);
            }
        }
    }
    events.sort_by(|a, b| a.ts.cmp(&b.ts).then_with(|| a.event_id.cmp(&b.event_id)));
    Ok(events)
}

pub fn read_events_from_file(project_dir: &Path, file: &str) -> Result<Vec<TraceEvent>> {
    let path = project_dir.join(".sdd").join(file);
    if !path.exists() {
        return Ok(Vec::new());
    }
    let text = fs::read_to_string(&path).with_context(|| format!("reading {}", path.display()))?;
    let mut events = Vec::new();
    for (index, chunk) in json_object_chunks(&text).into_iter().enumerate() {
        if chunk.len() > 8_192 {
            continue;
        }
        if let Ok(value) = serde_json::from_str::<Value>(&chunk) {
            events.push(TraceEvent::from_payload(value, file, index + 1));
        }
    }
    Ok(events)
}

pub fn events_for_run(project_dir: &Path, run_id: &str) -> Result<Vec<TraceEvent>> {
    let events = read_events(project_dir)?;
    let mut children = BTreeMap::<String, Vec<String>>::new();
    for event in &events {
        if let Some(parent) = &event.parent_run_id {
            children
                .entry(parent.clone())
                .or_default()
                .push(event.run_id.clone());
        }
    }
    let mut wanted = BTreeSet::from([run_id.to_string()]);
    let mut queue = vec![run_id.to_string()];
    while let Some(current) = queue.pop() {
        if let Some(items) = children.get(&current) {
            for child in items {
                if wanted.insert(child.clone()) {
                    queue.push(child.clone());
                }
            }
        }
    }
    Ok(events
        .into_iter()
        .filter(|event| wanted.contains(&event.run_id))
        .collect())
}

pub fn trace_tree(project_dir: &Path, run_id: &str) -> Result<TraceTree> {
    let events = events_for_run(project_dir, run_id)?;
    Ok(build_tree(run_id, &events))
}

pub fn summary(events: &[TraceEvent]) -> TraceSummary {
    let mut out = TraceSummary {
        total_events: events.len(),
        ..TraceSummary::default()
    };
    let run_ids = events
        .iter()
        .map(|event| event.run_id.clone())
        .collect::<BTreeSet<_>>();
    for event in events {
        *out.by_kind.entry(event.kind.clone()).or_insert(0) += 1;
        if let Some(status) = &event.status {
            *out.by_status.entry(status.clone()).or_insert(0) += 1;
        }
        if let Some(stage) = &event.stage {
            *out.by_stage.entry(stage.clone()).or_insert(0) += 1;
        }
        if out
            .latest_ts
            .as_deref()
            .is_none_or(|ts| event.ts.as_str() > ts)
        {
            out.latest_ts = Some(event.ts.clone());
        }
    }
    let roots = events
        .iter()
        .filter(|event| {
            event
                .parent_run_id
                .as_ref()
                .is_none_or(|parent| !run_ids.contains(parent))
        })
        .map(|event| event.run_id.clone())
        .collect::<BTreeSet<_>>()
        .into_iter()
        .collect::<Vec<_>>();
    out.root_count = roots.len();
    out.roots = roots.into_iter().take(50).collect();
    out
}

pub fn doctor(project_dir: &Path) -> Result<TraceDoctorReport> {
    let mut report = TraceDoctorReport {
        status: "pass",
        files: BTreeMap::new(),
        warnings: Vec::new(),
    };
    for file in TRACE_FILES {
        match read_events_from_file(project_dir, file) {
            Ok(events) => {
                report.files.insert((*file).to_string(), events.len());
            }
            Err(error) => {
                report.status = "fail";
                report.warnings.push(error.to_string());
            }
        }
    }
    Ok(report)
}

pub fn filter_by_orchestration(events: Vec<TraceEvent>, orchestration: &str) -> Vec<TraceEvent> {
    let wanted = orchestration.trim();
    events
        .into_iter()
        .filter(|event| event.orchestration.as_deref() == Some(wanted))
        .collect()
}

pub fn validate_sdd_uri(uri: &str) -> Result<Vec<String>> {
    let Some(rest) = uri.strip_prefix("sdd://") else {
        anyhow::bail!("resource uri must start with sdd://");
    };
    let segments = rest
        .split('/')
        .map(str::trim)
        .filter(|item| !item.is_empty())
        .map(|item| {
            if item == "." || item == ".." || item.contains('\\') {
                anyhow::bail!("invalid sdd resource segment `{item}`");
            }
            Ok(item.to_string())
        })
        .collect::<Result<Vec<_>>>()?;
    Ok(segments)
}

fn build_tree(run_id: &str, events: &[TraceEvent]) -> TraceTree {
    let own = events
        .iter()
        .filter(|event| event.run_id == run_id)
        .cloned()
        .collect::<Vec<_>>();
    let children = events
        .iter()
        .filter(|event| event.parent_run_id.as_deref() == Some(run_id))
        .map(|event| event.run_id.clone())
        .collect::<BTreeSet<_>>();
    TraceTree {
        run_id: run_id.to_string(),
        events: own,
        children: children
            .into_iter()
            .map(|child| build_tree(&child, events))
            .collect(),
    }
}

fn string_field(map: Option<&Map<String, Value>>, key: &str) -> Option<String> {
    map.and_then(|object| object.get(key))
        .and_then(|value| match value {
            Value::String(text) if !text.trim().is_empty() => Some(redact_trace_text(text)),
            Value::Number(_) | Value::Bool(_) => Some(value.to_string()),
            _ => None,
        })
}

fn u128_field(map: Option<&Map<String, Value>>, key: &str) -> Option<u128> {
    map.and_then(|object| object.get(key))
        .and_then(Value::as_u64)
        .map(u128::from)
}

fn redact_value(value: Value) -> Value {
    match value {
        Value::String(text) => Value::String(redact_trace_text(&text)),
        Value::Array(items) => Value::Array(items.into_iter().map(redact_value).collect()),
        Value::Object(map) => Value::Object(
            map.into_iter()
                .map(|(key, value)| (key, redact_value(value)))
                .collect(),
        ),
        other => other,
    }
}

fn legacy_run_id(source_file: &str, source_line: usize, payload: &Value) -> String {
    make_run_id(&format!("{source_file}:{source_line}:{payload}"))
}

fn event_id(source_file: &str, source_line: usize, payload: &Value) -> String {
    make_event_id(&format!("{source_file}:{source_line}:{payload}"))
}

fn next_line(project_dir: &Path, file: &str) -> usize {
    let path: PathBuf = project_dir.join(".sdd").join(file);
    fs::read_to_string(path)
        .map(|text| text.lines().count() + 1)
        .unwrap_or(1)
}

fn hex_prefix(bytes: &[u8], chars: usize) -> String {
    bytes
        .iter()
        .flat_map(|byte| {
            let high = byte >> 4;
            let low = byte & 0x0f;
            [hex_digit(high), hex_digit(low)]
        })
        .take(chars)
        .collect()
}

fn hex_digit(value: u8) -> char {
    match value {
        0..=9 => (b'0' + value) as char,
        _ => (b'a' + value - 10) as char,
    }
}

fn redact_trace_text(text: &str) -> String {
    let truncated = text.chars().take(256).collect::<String>();
    let lower = truncated.to_ascii_lowercase();
    let secret_markers = [
        "api_key",
        "apikey",
        "access_token",
        "authorization: bearer",
        "authorization=bearer",
        "password",
        "credential",
        "secret",
        "sk-",
        "sk_ant",
        "sk-ant",
        "xoxb-",
        "xoxp-",
        "ghp_",
        "github_pat_",
        "glpat-",
        "akia",
        "eyj",
    ];
    if secret_markers.iter().any(|marker| lower.contains(marker)) {
        "[REDACTED]".to_string()
    } else {
        truncated
    }
}

fn json_object_chunks(text: &str) -> Vec<String> {
    let mut chunks = Vec::new();
    for line in text.lines().map(str::trim).filter(|line| !line.is_empty()) {
        let parts = line.split("}{").collect::<Vec<_>>();
        for (index, part) in parts.iter().enumerate() {
            let mut chunk = (*part).to_string();
            if index > 0 && !chunk.starts_with('{') {
                chunk.insert(0, '{');
            }
            if index + 1 < parts.len() && !chunk.ends_with('}') {
                chunk.push('}');
            }
            if chunk.starts_with('{') && chunk.ends_with('}') && chunk.len() <= 8_192 {
                chunks.push(chunk);
            }
        }
    }
    chunks
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn normalizes_legacy_events_and_redacts_secrets() {
        let event = TraceEvent::from_payload(
            json!({
                "ts": "2026-06-11T20:00:00Z",
                "event": "PostToolUse",
                "tool_name": "Bash",
                "tool_input": {"command": "echo OPENAI_API_KEY=sk-secret-value"}
            }),
            "events.jsonl",
            1,
        );

        assert_eq!(event.schema_version, TRACE_SCHEMA_VERSION);
        assert_eq!(event.kind, "PostToolUse");
        assert_eq!(event.tool.as_deref(), Some("Bash"));
        assert!(!event.metadata.to_string().contains("sk-secret-value"));
    }

    #[test]
    fn reads_run_tree_with_children() {
        let dir = tempdir().unwrap();
        let root = dir.path();
        let parent = TraceEvent::from_payload(
            json!({"ts": "1", "run_id": "root", "kind": "stage", "status": "started"}),
            "runs.jsonl",
            1,
        );
        let child = TraceEvent::from_payload(
            json!({"ts": "2", "run_id": "child", "parent_run_id": "root", "kind": "tool"}),
            "runs.jsonl",
            2,
        );
        append_event(root, "runs.jsonl", &parent).unwrap();
        append_event(root, "runs.jsonl", &child).unwrap();

        let tree = trace_tree(root, "root").unwrap();
        assert_eq!(tree.run_id, "root");
        assert_eq!(tree.children.len(), 1);
        assert_eq!(tree.children[0].run_id, "child");
    }

    #[test]
    fn keeps_allowed_gate_reason_out_of_error() {
        let allowed = TraceEvent::from_payload(
            json!({
                "ts": "1",
                "kind": "task_gate",
                "status": "allowed",
                "reason": "quality gate passed"
            }),
            "task-gates.jsonl",
            1,
        );
        let blocked = TraceEvent::from_payload(
            json!({
                "ts": "2",
                "kind": "task_gate",
                "status": "blocked",
                "reason": "missing evidence"
            }),
            "task-gates.jsonl",
            2,
        );

        assert_eq!(allowed.error, None);
        assert_eq!(
            allowed.metadata.get("reason").and_then(Value::as_str),
            Some("quality gate passed")
        );
        assert_eq!(blocked.error.as_deref(), Some("missing evidence"));
    }

    #[test]
    fn rejects_path_traversal_in_sdd_uris() {
        assert!(validate_sdd_uri("sdd://artifact/demo/prd").is_ok());
        assert!(validate_sdd_uri("sdd://artifact/../prd").is_err());
        assert!(validate_sdd_uri("file:///tmp/nope").is_err());
    }

    #[test]
    fn parses_concatenated_legacy_json_objects() {
        let chunks = json_object_chunks(r#"{"a":1}{"b":"x\"y"}"#);
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0], r#"{"a":1}"#);
        assert_eq!(chunks[1], r#"{"b":"x\"y"}"#);
    }
}
