//! Cost tracking (issue #49): collect token usage exposed by agents, price
//! it against the configured model sheet, and aggregate per agent/session.
//! No runtime dependency is added — everything is std + the existing config.

use crate::config::{self, CostModel};
use crate::events::{self, Event, EventKind};
use anyhow::Result;
use serde::Serialize;
use std::collections::BTreeMap;
use std::path::PathBuf;

/// Token usage collected from an agent.
#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize)]
pub struct Usage {
    pub tokens_in: u64,
    pub tokens_out: u64,
}

/// One aggregated cost row (per agent, or global totals).
#[derive(Debug, Clone, Default, Serialize)]
pub struct CostStats {
    pub agent: String,
    pub sessions: u64,
    pub tokens_in: u64,
    pub tokens_out: u64,
    pub cost_usd: f64,
}

/// Price the usage against a model of the sheet (issue #49).
/// Unknown models fall back to the first model of the sheet.
pub fn estimate_cost(
    model: Option<&str>,
    usage: Usage,
    sheet: &BTreeMap<String, CostModel>,
) -> f64 {
    let m = model
        .and_then(|k| sheet.get(k))
        .or_else(|| sheet.values().next());
    match m {
        Some(cm) => {
            usage.tokens_in as f64 * cm.input_per_mtok / 1_000_000.0
                + usage.tokens_out as f64 * cm.output_per_mtok / 1_000_000.0
        }
        None => 0.0,
    }
}

/// Collect the usage an agent exposed for its latest session (issue #49).
///
/// Supported sources today:
/// - Claude Code: `~/.claude/projects/*/<session>.jsonl` — the last JSON
///   line carrying an `usage` object with `input_tokens`/`output_tokens`.
///
/// Returns None when the agent exposes nothing (safe, silent).
pub fn collect_usage(agent: &str, cwd: Option<&str>) -> Option<Usage> {
    match agent {
        "claude-code" => claude_usage(cwd),
        _ => None,
    }
}

/// Claude Code stores one JSONL per session under ~/.claude/projects/<enc>/.
/// The final line with a "usage" object holds the cumulative token counts.
fn claude_usage(cwd: Option<&str>) -> Option<Usage> {
    let root = config::home_dir()?.join(".claude/projects");
    let mut candidates: Vec<(String, PathBuf)> = Vec::new();
    for entry in std::fs::read_dir(&root).ok()?.flatten() {
        let dir = entry.path();
        if !dir.is_dir() {
            continue;
        }
        for f in std::fs::read_dir(&dir).ok()?.flatten() {
            let p = f.path();
            if p.extension().and_then(|e| e.to_str()) == Some("jsonl") {
                if let Ok(mtime) = f.metadata().and_then(|m| m.modified()) {
                    candidates.push((mtime_key(mtime), p));
                }
            }
        }
    }
    candidates.sort();
    let last = candidates.last()?.1.clone();
    parse_claude_file(&last, cwd)
}

fn mtime_key(t: std::time::SystemTime) -> String {
    format!(
        "{:?}",
        t.duration_since(std::time::UNIX_EPOCH).unwrap_or_default()
    )
}

/// Parse a Claude Code session file: the last line whose "usage" object has
/// input/output token counts. Filters by cwd when given (the file embeds it
/// as `cwd` on summary lines).
fn parse_claude_file(path: &std::path::Path, cwd: Option<&str>) -> Option<Usage> {
    let text = std::fs::read_to_string(path).ok()?;
    let mut usage = None;
    for line in text.lines() {
        let v: serde_json::Value = serde_json::from_str(line).ok()?;
        if v.get("usage").is_none() {
            continue;
        }
        if let Some(want) = cwd {
            if v.get("cwd").and_then(|c| c.as_str()) != Some(want) {
                continue;
            }
        }
        let u = v.get("usage")?;
        let tin = u.get("input_tokens").and_then(|n| n.as_u64())?;
        let tout = u.get("output_tokens").and_then(|n| n.as_u64())?;
        usage = Some(Usage {
            tokens_in: tin,
            tokens_out: tout,
        });
    }
    usage
}

/// Record a session's usage: persist it on the session record and emit a
/// `cost` event (axe 8 journal). Called when an agent session finishes.
pub fn record_session_usage(
    app: &crate::app::App,
    agent: &str,
    session_id: &str,
    cwd: Option<&str>,
) -> Result<()> {
    let Some(usage) = collect_usage(agent, cwd) else {
        return Ok(());
    };
    let model = app
        .catalog
        .resolve(agent)
        .and_then(|a| a.cost_model.as_deref());
    let sheet = app
        .config
        .settings
        .cost_models
        .clone()
        .unwrap_or_else(config::default_cost_models);
    let cost = estimate_cost(model, usage, &sheet);

    // Persist on the session record.
    let mut sf = crate::sessions::load(app)?;
    if let Some(r) = sf.sessions.get_mut(session_id) {
        r.tokens_in = Some(usage.tokens_in);
        r.tokens_out = Some(usage.tokens_out);
        r.cost_usd = Some(cost);
        crate::sessions::save(app, &sf)?;
    }

    // Journal a dedicated event (stable --json contract, additive fields).
    app.emit(
        &Event::now(EventKind::Cost)
            .with_agent(agent.to_string())
            .with_session(session_id)
            .with_usage(usage.tokens_in, usage.tokens_out)
            .with_cost(cost),
    );
    Ok(())
}

/// Aggregate cost events into per-agent rows (issue #49). `--json` output
/// stays stable: agent, sessions, tokens_in, tokens_out, cost_usd.
pub fn aggregate(
    app: &crate::app::App,
    period: Option<&str>,
) -> Result<Vec<CostStats>> {
    let cut = crate::commands::stats_cmd::cutoff(period)?;
    let events = events::read_events(&app.events_dir(), 0);
    let mut map: BTreeMap<String, CostStats> = BTreeMap::new();
    for ev in events {
        if ev.kind != EventKind::Cost {
            continue;
        }
        if let Some(cut) = &cut {
            if let Ok(ts) = chrono::DateTime::parse_from_rfc3339(&ev.ts) {
                if ts.with_timezone(&chrono::Utc) < *cut {
                    continue;
                }
            }
        }
        let agent = ev.agent.clone().unwrap_or_else(|| "unknown".to_string());
        let row = map.entry(agent.clone()).or_default();
        row.agent = agent;
        row.sessions += 1;
        row.tokens_in += ev.tokens_in.unwrap_or(0);
        row.tokens_out += ev.tokens_out.unwrap_or(0);
        row.cost_usd += ev.cost_usd.unwrap_or(0.0);
    }
    Ok(map.into_values().collect())
}

/// Total across every row (used by the weekly/monthly report).
pub fn totals(rows: &[CostStats]) -> CostStats {
    let mut t = CostStats {
        agent: "total".to_string(),
        ..Default::default()
    };
    for r in rows {
        t.sessions += r.sessions;
        t.tokens_in += r.tokens_in;
        t.tokens_out += r.tokens_out;
        t.cost_usd += r.cost_usd;
    }
    t
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn estimate_cost_uses_model_prices() {
        let sheet = config::default_cost_models();
        // 1M input @ $3 + 500k output @ $15 = 3 + 7.5 = $10.5
        let u = Usage {
            tokens_in: 1_000_000,
            tokens_out: 500_000,
        };
        let c = estimate_cost(Some("claude-sonnet-4"), u, &sheet);
        assert!((c - 10.5).abs() < 1e-9, "got {c}");
    }

    #[test]
    fn estimate_cost_falls_back_to_first_model() {
        let sheet = config::default_cost_models();
        let c = estimate_cost(Some("nope-model"), Usage { tokens_in: 0, tokens_out: 0 }, &sheet);
        assert_eq!(c, 0.0);
    }

    #[test]
    fn estimate_cost_unknown_sheet_is_zero() {
        let sheet = BTreeMap::new();
        let c = estimate_cost(None, Usage { tokens_in: 1000, tokens_out: 1000 }, &sheet);
        assert_eq!(c, 0.0);
    }

    #[test]
    fn parse_claude_file_extracts_last_usage() {
        let dir = std::env::temp_dir().join(format!("am-costs-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let f = dir.join("session.jsonl");
        std::fs::write(
            &f,
            concat!(
                "{\"type\":\"assistant\",\"message\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":5}},\"cwd\":\"/x\"}\n",
                "{\"type\":\"summary\",\"usage\":{\"input_tokens\":100,\"output_tokens\":40},\"cwd\":\"/x\"}\n",
            ),
        )
        .unwrap();
        let u = parse_claude_file(&f, Some("/x")).unwrap();
        assert_eq!(u.tokens_in, 100);
        assert_eq!(u.tokens_out, 40);
        // Wrong cwd → filtered out.
        assert!(parse_claude_file(&f, Some("/other")).is_none());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn aggregate_sums_per_agent() {
        let app = test_app("costs-agg");
        app.emit(
            &Event::now(EventKind::Cost)
                .with_agent("a".to_string())
                .with_usage(10, 5)
                .with_cost(0.5),
        );
        app.emit(
            &Event::now(EventKind::Cost)
                .with_agent("a".to_string())
                .with_usage(20, 10)
                .with_cost(1.0),
        );
        app.emit(
            &Event::now(EventKind::Cost)
                .with_agent("b".to_string())
                .with_usage(100, 50)
                .with_cost(2.0),
        );
        let rows = aggregate(&app, None).unwrap();
        assert_eq!(rows.len(), 2);
        let a = rows.iter().find(|r| r.agent == "a").unwrap();
        assert_eq!(a.sessions, 2);
        assert_eq!(a.tokens_in, 30);
        assert_eq!(a.tokens_out, 15);
        assert!((a.cost_usd - 1.5).abs() < 1e-9);
        let t = totals(&rows);
        assert!((t.cost_usd - 3.5).abs() < 1e-9);
        assert_eq!(t.sessions, 3);
    }

    fn test_app(tag: &str) -> crate::app::App {
        let dir = tempfile::tempdir().unwrap();
        let cfg = dir.path().join("config.yaml");
        std::fs::write(
            &cfg,
            concat!(
                "version: \"1.0\"\n",
                "settings:\n",
                "  auto_install_deps: false\n",
                "  confirm_before_run: false\n",
                "agents: []\n",
            ),
        )
        .unwrap();
        use clap::Parser;
        let cli = crate::cli::Cli::parse_from([
            "am",
            "--config",
            cfg.to_str().unwrap(),
            "list",
        ]);
        let mut app = crate::app::App::from_cli(cli).expect("app should build");
        app.paths.state_file = dir.path().join(format!("state-{tag}.json"));
        app.paths.log_dir = dir.path().join("logs");
        app
    }
}