//! In-memory CostRollup register, keyed by project root.
//!
//! Keying by `project_root` provides test isolation: each test's `TempDir`
//! maps to its own slice of the store, so parallel tests cannot pollute
//! each other's data.
//!
//! Populated on server startup from existing JSONL records
//! ([`init_from_disk`]) and kept current by the cost-rollup subscriber
//! (`agents::pool::cost_rollup_subscriber`) which fires on every terminal
//! pipeline stage transition.
//!
//! All readers call [`get_rollup`] or [`all_rollups`] instead of
//! re-walking `token_usage.jsonl`.

use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{OnceLock, RwLock};

use super::token::{AgentTokenCost, TokenCostSummary};

// ── Types ─────────────────────────────────────────────────────────────────────

/// Snapshotted cost totals for a story at the time it reached a terminal stage.
#[derive(Debug, Clone)]
pub struct CostRollup {
    /// The story this rollup belongs to.
    pub story_id: String,
    /// Total USD spend across all agents.
    pub total_cost_usd: f64,
    /// Per-agent breakdown, sorted alphabetically by agent name.
    pub agents: Vec<AgentTokenCost>,
    /// When the rollup was written (terminal transition timestamp or, for
    /// startup-bootstrapped entries, the latest token record timestamp).
    pub recorded_at: DateTime<Utc>,
}

impl CostRollup {
    /// Convert to a [`TokenCostSummary`] for callers that expect the legacy type.
    pub fn as_summary(&self) -> TokenCostSummary {
        TokenCostSummary {
            total_cost_usd: self.total_cost_usd,
            agents: self.agents.clone(),
        }
    }
}

// ── Global store ──────────────────────────────────────────────────────────────

/// Composite key: (canonical project root, story_id).
type StoreKey = (PathBuf, String);

static STORE: OnceLock<RwLock<HashMap<StoreKey, CostRollup>>> = OnceLock::new();

fn store() -> &'static RwLock<HashMap<StoreKey, CostRollup>> {
    STORE.get_or_init(|| RwLock::new(HashMap::new()))
}

fn canonical(p: &Path) -> PathBuf {
    p.canonicalize().unwrap_or_else(|_| p.to_path_buf())
}

/// Look up the cost rollup for `story_id` within `project_root`.
///
/// Returns `None` for stories that have not yet reached a terminal stage
/// (or have never had any token records).
pub fn get_rollup(project_root: &Path, story_id: &str) -> Option<CostRollup> {
    let key = (canonical(project_root), story_id.to_string());
    store().read().ok()?.get(&key).cloned()
}

/// Write or overwrite the cost rollup for `story_id` within `project_root`.
pub fn set_rollup(project_root: &Path, story_id: impl Into<String>, rollup: CostRollup) {
    let key = (canonical(project_root), story_id.into());
    if let Ok(mut guard) = store().write() {
        guard.insert(key, rollup);
    }
}

/// Return a snapshot of all rollups for `project_root`, in arbitrary order.
pub fn all_rollups(project_root: &Path) -> Vec<CostRollup> {
    let root = canonical(project_root);
    store()
        .read()
        .map(|g| {
            g.iter()
                .filter(|((r, _), _)| r == &root)
                .map(|(_, v)| v.clone())
                .collect()
        })
        .unwrap_or_default()
}

// ── Startup bootstrap ─────────────────────────────────────────────────────────

/// Pre-populate the register from `token_usage.jsonl`.
///
/// Called once at server startup before the live subscriber starts listening.
/// Stories with zero cost are skipped.  `recorded_at` is set to the latest
/// record timestamp for each story so the 24 h window in the cost command
/// reflects when work was actually done rather than when the server restarted.
pub fn init_from_disk(project_root: &Path) {
    let records = match crate::agents::token_usage::read_all(project_root) {
        Ok(r) => r,
        Err(e) => {
            crate::slog_warn!("[cost-rollup] init_from_disk failed to read token records: {e}");
            return;
        }
    };

    if records.is_empty() {
        return;
    }

    // Collect unique story IDs.
    let mut story_ids: Vec<&str> = records.iter().map(|r| r.story_id.as_str()).collect();
    story_ids.sort_unstable();
    story_ids.dedup();

    for sid in story_ids {
        let summary = super::token::aggregate_for_story(&records, sid);
        if summary.total_cost_usd == 0.0 && summary.agents.is_empty() {
            continue;
        }

        // Use the latest record timestamp as recorded_at.
        let recorded_at = records
            .iter()
            .filter(|r| r.story_id == sid)
            .filter_map(|r| chrono::DateTime::parse_from_rfc3339(&r.timestamp).ok())
            .map(|ts| ts.with_timezone(&Utc))
            .max()
            .unwrap_or_else(Utc::now);

        set_rollup(
            project_root,
            sid,
            CostRollup {
                story_id: sid.to_string(),
                total_cost_usd: summary.total_cost_usd,
                agents: summary.agents,
                recorded_at,
            },
        );
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agents::TokenUsage;
    use crate::agents::token_usage::TokenUsageRecord;

    fn dummy_root() -> PathBuf {
        std::env::temp_dir().join("cost_rollup_dummy_root")
    }

    fn make_record(story_id: &str, agent: &str, cost: f64, ts: &str) -> TokenUsageRecord {
        TokenUsageRecord {
            story_id: story_id.to_string(),
            agent_name: agent.to_string(),
            timestamp: ts.to_string(),
            model: None,
            usage: TokenUsage {
                input_tokens: 100,
                output_tokens: 50,
                cache_creation_input_tokens: 0,
                cache_read_input_tokens: 0,
                total_cost_usd: cost,
            },
        }
    }

    #[test]
    fn get_rollup_returns_none_for_unknown_story() {
        let root = dummy_root();
        assert!(get_rollup(&root, "999_no_such_story_xyzzy").is_none());
    }

    #[test]
    fn set_and_get_rollup_roundtrip() {
        let root = dummy_root();
        let rollup = CostRollup {
            story_id: "1017_test_roundtrip".to_string(),
            total_cost_usd: 3.50,
            agents: vec![],
            recorded_at: Utc::now(),
        };
        set_rollup(&root, "1017_test_roundtrip", rollup);
        let got = get_rollup(&root, "1017_test_roundtrip").expect("rollup must exist after set");
        assert!((got.total_cost_usd - 3.50).abs() < f64::EPSILON);
    }

    #[test]
    fn init_from_disk_populates_store() {
        let tmp = tempfile::tempdir().unwrap();
        let records = vec![
            make_record("1017_init_a", "coder-1", 1.0, "2026-01-01T00:00:00Z"),
            make_record("1017_init_a", "qa-1", 0.5, "2026-01-02T00:00:00Z"),
            make_record("1017_init_b", "coder-1", 2.0, "2026-01-01T12:00:00Z"),
        ];
        for r in &records {
            crate::agents::token_usage::append_record(tmp.path(), r).unwrap();
        }
        init_from_disk(tmp.path());

        let a = get_rollup(tmp.path(), "1017_init_a").expect("story a must be in store");
        assert!((a.total_cost_usd - 1.5).abs() < f64::EPSILON);
        assert_eq!(a.agents.len(), 2);

        let b = get_rollup(tmp.path(), "1017_init_b").expect("story b must be in store");
        assert!((b.total_cost_usd - 2.0).abs() < f64::EPSILON);
    }

    #[test]
    fn init_from_disk_uses_latest_timestamp_as_recorded_at() {
        let tmp = tempfile::tempdir().unwrap();
        let records = vec![
            make_record("1017_ts_test", "coder-1", 1.0, "2026-01-01T00:00:00Z"),
            make_record("1017_ts_test", "qa-1", 0.5, "2026-01-03T00:00:00Z"),
        ];
        for r in &records {
            crate::agents::token_usage::append_record(tmp.path(), r).unwrap();
        }
        init_from_disk(tmp.path());

        let rollup = get_rollup(tmp.path(), "1017_ts_test").expect("rollup must exist");
        // recorded_at should be 2026-01-03 (the later timestamp)
        assert_eq!(rollup.recorded_at.date_naive().to_string(), "2026-01-03");
    }

    #[test]
    fn all_rollups_isolated_by_project_root() {
        let tmp_a = tempfile::tempdir().unwrap();
        let tmp_b = tempfile::tempdir().unwrap();
        set_rollup(
            tmp_a.path(),
            "1017_all_a",
            CostRollup {
                story_id: "1017_all_a".to_string(),
                total_cost_usd: 1.0,
                agents: vec![],
                recorded_at: Utc::now(),
            },
        );
        set_rollup(
            tmp_b.path(),
            "1017_all_b",
            CostRollup {
                story_id: "1017_all_b".to_string(),
                total_cost_usd: 2.0,
                agents: vec![],
                recorded_at: Utc::now(),
            },
        );
        let all_a = all_rollups(tmp_a.path());
        assert!(all_a.iter().any(|r| r.story_id == "1017_all_a"));
        assert!(!all_a.iter().any(|r| r.story_id == "1017_all_b"));

        let all_b = all_rollups(tmp_b.path());
        assert!(all_b.iter().any(|r| r.story_id == "1017_all_b"));
        assert!(!all_b.iter().any(|r| r.story_id == "1017_all_a"));
    }
}
