use std::collections::{BTreeMap, HashSet};

use serde::Serialize;
use time::{Date, Duration};

use crate::lfd::sessions::types::{
    ContextSnapshot, PersistedSessionEvent, Session, SessionEvent, TurnUsage,
};

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
pub struct TokenTotals {
    pub input: u64,
    pub output: u64,
    pub reasoning: u64,
    pub cache_read: u64,
    pub cache_write: u64,
}

impl TokenTotals {
    fn add_turn_usage(&mut self, usage: &TurnUsage) {
        self.input += usage.input_tokens;
        self.output += usage.output_tokens;
        self.reasoning += usage.reasoning_tokens.unwrap_or(0);
        self.cache_read += usage.cache_read_tokens.unwrap_or(0);
        self.cache_write += usage.cache_write_tokens.unwrap_or(0);
    }

    fn merge(&mut self, other: &TokenTotals) {
        self.input += other.input;
        self.output += other.output;
        self.reasoning += other.reasoning;
        self.cache_read += other.cache_read;
        self.cache_write += other.cache_write;
    }
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SessionUsageAggregate {
    pub tokens: TokenTotals,
    pub turns: u64,
    pub context: Option<ContextSnapshot>,
    pub models: BTreeMap<String, u64>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
pub struct StepUsageAggregate {
    #[serde(flatten)]
    pub tokens: TokenTotals,
    pub sessions: u64,
    pub turns: u64,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct WaveUsageAggregate {
    pub tokens: TokenTotals,
    pub sessions: u64,
    pub turns: u64,
    pub models: BTreeMap<String, u64>,
    pub by_step: BTreeMap<String, StepUsageAggregate>,
}

#[derive(Debug, Clone)]
pub struct UsageSessionData {
    pub session: Session,
    pub events: Vec<PersistedSessionEvent>,
    pub wave_id: Option<String>,
    pub flow: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GroupBy {
    Wave,
    Flow,
    Step,
    Model,
    Source,
}

impl GroupBy {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Wave => "wave",
            Self::Flow => "flow",
            Self::Step => "step",
            Self::Model => "model",
            Self::Source => "source",
        }
    }
}

impl std::str::FromStr for GroupBy {
    type Err = InvalidGroupBy;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "wave" => Ok(Self::Wave),
            "flow" => Ok(Self::Flow),
            "step" => Ok(Self::Step),
            "model" => Ok(Self::Model),
            "source" => Ok(Self::Source),
            _ => Err(InvalidGroupBy),
        }
    }
}

#[derive(Debug)]
pub struct InvalidGroupBy;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimeBucket {
    Day,
    Week,
    Month,
}

impl TimeBucket {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Day => "day",
            Self::Week => "week",
            Self::Month => "month",
        }
    }

    fn start_date(self, created_at: time::OffsetDateTime) -> Date {
        let date = created_at.date();
        match self {
            Self::Day => date,
            Self::Week => {
                let days_from_monday = i64::from(date.weekday().number_days_from_monday());
                date - Duration::days(days_from_monday)
            }
            Self::Month => Date::from_calendar_date(date.year(), date.month(), 1)
                .expect("calendar month start should always be valid"),
        }
    }

    fn format_period(self, date: Date) -> String {
        match self {
            Self::Day | Self::Week => date.to_string(),
            Self::Month => format!("{}-{:02}", date.year(), date.month() as u8),
        }
    }
}

impl std::str::FromStr for TimeBucket {
    type Err = InvalidTimeBucket;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "day" => Ok(Self::Day),
            "week" => Ok(Self::Week),
            "month" => Ok(Self::Month),
            _ => Err(InvalidTimeBucket),
        }
    }
}

#[derive(Debug)]
pub struct InvalidTimeBucket;

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct UsageSummaryGroupAggregate {
    pub key: String,
    pub tokens: TokenTotals,
    pub sessions: u64,
    pub turns: u64,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct UsageTimeseriesBucketAggregate {
    pub period: String,
    pub groups: Vec<UsageSummaryGroupAggregate>,
}

pub fn aggregate_wave_usage(
    sessions: &[(Session, Vec<PersistedSessionEvent>)],
) -> WaveUsageAggregate {
    let mut aggregate = WaveUsageAggregate::default();

    for (session, events) in sessions {
        let session_usage = aggregate_session_events(events, None);

        aggregate.sessions += 1;
        aggregate.turns += session_usage.turns;
        aggregate.tokens.merge(&session_usage.tokens);
        merge_model_counts(&mut aggregate.models, &session_usage.models);

        let step = step_key(session);
        let step_usage = aggregate.by_step.entry(step).or_default();
        step_usage.sessions += 1;
        step_usage.turns += session_usage.turns;
        step_usage.tokens.merge(&session_usage.tokens);
    }

    aggregate
}

pub fn aggregate_summary(
    group_by: GroupBy,
    sessions: &[UsageSessionData],
    model_filter: Option<&str>,
) -> Vec<UsageSummaryGroupAggregate> {
    match group_by {
        GroupBy::Wave => aggregate_summary_by_key(sessions, model_filter, |data| {
            data.wave_id
                .clone()
                .or_else(|| data.session.config.wave.clone())
                .unwrap_or_else(|| "unknown".to_string())
        }),
        GroupBy::Flow => aggregate_summary_by_key(sessions, model_filter, |data| {
            data.flow.clone().unwrap_or_else(|| "unknown".to_string())
        }),
        GroupBy::Step => {
            aggregate_summary_by_key(sessions, model_filter, |data| step_key(&data.session))
        }
        GroupBy::Model => aggregate_summary_by_model(sessions, model_filter),
        GroupBy::Source => aggregate_summary_by_source(sessions),
    }
}

pub fn aggregate_timeseries(
    bucket: TimeBucket,
    group_by: GroupBy,
    sessions: Vec<UsageSessionData>,
    model_filter: Option<&str>,
) -> Vec<UsageTimeseriesBucketAggregate> {
    let mut sessions_by_period: BTreeMap<Date, Vec<UsageSessionData>> = BTreeMap::new();
    for data in sessions {
        let period_start = bucket.start_date(data.session.created_at);
        sessions_by_period
            .entry(period_start)
            .or_default()
            .push(data);
    }

    sessions_by_period
        .into_iter()
        .map(
            |(period_start, grouped_sessions)| UsageTimeseriesBucketAggregate {
                period: bucket.format_period(period_start),
                groups: aggregate_summary(group_by, &grouped_sessions, model_filter),
            },
        )
        .collect()
}

pub fn aggregate_session_events(
    events: &[PersistedSessionEvent],
    model_filter: Option<&str>,
) -> SessionUsageAggregate {
    let mut aggregate = SessionUsageAggregate::default();

    for event in events {
        match &event.event {
            SessionEvent::TurnUsage { usage, .. } => {
                if let Some(filter) = model_filter {
                    if usage.model.as_deref() != Some(filter) {
                        continue;
                    }
                }

                aggregate.turns += 1;
                aggregate.tokens.add_turn_usage(usage);
                if let Some(model) = usage.model.as_ref() {
                    *aggregate.models.entry(model.clone()).or_insert(0) += 1;
                }
            }
            SessionEvent::ContextSnapshot { snapshot } if aggregate.context.is_none() => {
                aggregate.context = Some(snapshot.clone());
            }
            _ => {}
        }
    }

    aggregate
}

fn aggregate_summary_by_key(
    sessions: &[UsageSessionData],
    model_filter: Option<&str>,
    key_fn: impl Fn(&UsageSessionData) -> String,
) -> Vec<UsageSummaryGroupAggregate> {
    let mut groups: BTreeMap<String, UsageSummaryGroupAggregate> = BTreeMap::new();

    for data in sessions {
        let usage = aggregate_session_events(&data.events, model_filter);
        if model_filter.is_some() && usage.turns == 0 {
            continue;
        }

        let key = key_fn(data);
        let group = groups
            .entry(key.clone())
            .or_insert(UsageSummaryGroupAggregate {
                key,
                tokens: TokenTotals::default(),
                sessions: 0,
                turns: 0,
            });
        group.tokens.merge(&usage.tokens);
        group.sessions += 1;
        group.turns += usage.turns;
    }

    groups.into_values().collect()
}

fn aggregate_summary_by_model(
    sessions: &[UsageSessionData],
    model_filter: Option<&str>,
) -> Vec<UsageSummaryGroupAggregate> {
    let mut groups: BTreeMap<String, UsageSummaryGroupAggregate> = BTreeMap::new();

    for data in sessions {
        let mut models_seen = HashSet::new();

        for event in &data.events {
            let SessionEvent::TurnUsage { usage, .. } = &event.event else {
                continue;
            };
            let Some(model) = usage.model.as_deref() else {
                continue;
            };
            if let Some(filter) = model_filter {
                if model != filter {
                    continue;
                }
            }

            let group = groups
                .entry(model.to_string())
                .or_insert(UsageSummaryGroupAggregate {
                    key: model.to_string(),
                    tokens: TokenTotals::default(),
                    sessions: 0,
                    turns: 0,
                });
            group.tokens.add_turn_usage(usage);
            group.turns += 1;
            if models_seen.insert(model.to_string()) {
                group.sessions += 1;
            }
        }
    }

    groups.into_values().collect()
}

fn aggregate_summary_by_source(sessions: &[UsageSessionData]) -> Vec<UsageSummaryGroupAggregate> {
    let mut groups: BTreeMap<String, UsageSummaryGroupAggregate> = BTreeMap::new();

    for data in sessions {
        let session_usage = aggregate_session_events(&data.events, None);
        let Some(snapshot) = session_usage.context else {
            continue;
        };

        for (source, source_tokens) in snapshot.sources {
            let group = groups
                .entry(source.clone())
                .or_insert(UsageSummaryGroupAggregate {
                    key: source,
                    tokens: TokenTotals::default(),
                    sessions: 0,
                    turns: 0,
                });
            group.tokens.input += source_tokens;
            group.sessions += 1;
        }
    }

    groups.into_values().collect()
}

fn step_key(session: &Session) -> String {
    let step = session.config.step.trim();
    if step.is_empty() {
        return "unknown".to_string();
    }
    step.to_string()
}

fn merge_model_counts(target: &mut BTreeMap<String, u64>, source: &BTreeMap<String, u64>) {
    for (model, count) in source {
        *target.entry(model.clone()).or_insert(0) += count;
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::lfd::id::LfdId;
    use crate::lfd::sessions::types::{SessionConfig, SessionStatus};
    use time::format_description::well_known::Rfc3339;
    use time::OffsetDateTime;

    #[test]
    fn aggregate_session_events_sums_turn_usage_and_context() {
        let events = vec![
            context_event(HashMap::from([
                ("step".to_string(), 100),
                ("diff".to_string(), 400),
            ])),
            turn_usage_event(TurnUsage {
                input_tokens: 120,
                output_tokens: 25,
                reasoning_tokens: Some(10),
                cache_read_tokens: Some(5),
                cache_write_tokens: Some(3),
                model: Some("claude-sonnet-4".to_string()),
                cost_usd: None,
            }),
            turn_usage_event(TurnUsage {
                input_tokens: 80,
                output_tokens: 20,
                reasoning_tokens: None,
                cache_read_tokens: None,
                cache_write_tokens: None,
                model: Some("claude-sonnet-4".to_string()),
                cost_usd: None,
            }),
        ];

        let usage = aggregate_session_events(&events, None);
        assert_eq!(usage.tokens.input, 200);
        assert_eq!(usage.tokens.output, 45);
        assert_eq!(usage.tokens.reasoning, 10);
        assert_eq!(usage.tokens.cache_read, 5);
        assert_eq!(usage.tokens.cache_write, 3);
        assert_eq!(usage.turns, 2);
        assert_eq!(usage.models.get("claude-sonnet-4"), Some(&2));
        assert_eq!(
            usage
                .context
                .expect("context snapshot should exist")
                .sources
                .get("diff"),
            Some(&400)
        );
    }

    #[test]
    fn aggregate_wave_usage_rolls_up_step_breakdown() {
        let session_a = test_session("implement");
        let session_b = test_session("gate");
        let wave_usage = aggregate_wave_usage(&[
            (
                session_a,
                vec![turn_usage_event(TurnUsage {
                    input_tokens: 150,
                    output_tokens: 30,
                    reasoning_tokens: Some(5),
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                })],
            ),
            (
                session_b,
                vec![turn_usage_event(TurnUsage {
                    input_tokens: 90,
                    output_tokens: 12,
                    reasoning_tokens: None,
                    cache_read_tokens: Some(8),
                    cache_write_tokens: Some(4),
                    model: Some("claude-haiku-4-5".to_string()),
                    cost_usd: None,
                })],
            ),
        ]);

        assert_eq!(wave_usage.sessions, 2);
        assert_eq!(wave_usage.turns, 2);
        assert_eq!(wave_usage.tokens.input, 240);
        assert_eq!(wave_usage.models.get("claude-sonnet-4"), Some(&1));
        assert_eq!(wave_usage.models.get("claude-haiku-4-5"), Some(&1));

        let implement = wave_usage.by_step.get("implement").expect("implement step");
        assert_eq!(implement.sessions, 1);
        assert_eq!(implement.tokens.input, 150);

        let gate = wave_usage.by_step.get("gate").expect("gate step");
        assert_eq!(gate.sessions, 1);
        assert_eq!(gate.tokens.input, 90);
    }

    #[test]
    fn aggregate_summary_by_model_honors_model_filter() {
        let records = vec![UsageSessionData {
            session: test_session("implement"),
            events: vec![
                turn_usage_event(TurnUsage {
                    input_tokens: 100,
                    output_tokens: 20,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                }),
                turn_usage_event(TurnUsage {
                    input_tokens: 60,
                    output_tokens: 8,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-haiku-4-5".to_string()),
                    cost_usd: None,
                }),
            ],
            wave_id: Some("wave-alpha".to_string()),
            flow: Some("build".to_string()),
        }];

        let groups = aggregate_summary(GroupBy::Model, &records, Some("claude-haiku-4-5"));
        assert_eq!(groups.len(), 1);
        assert_eq!(groups[0].key, "claude-haiku-4-5");
        assert_eq!(groups[0].tokens.input, 60);
        assert_eq!(groups[0].sessions, 1);
        assert_eq!(groups[0].turns, 1);
    }

    #[test]
    fn aggregate_summary_by_source_uses_context_snapshot() {
        let records = vec![UsageSessionData {
            session: test_session("implement"),
            events: vec![
                context_event(HashMap::from([
                    ("step".to_string(), 120),
                    ("diff".to_string(), 340),
                ])),
                turn_usage_event(TurnUsage {
                    input_tokens: 200,
                    output_tokens: 40,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                }),
            ],
            wave_id: Some("wave-alpha".to_string()),
            flow: Some("build".to_string()),
        }];

        let groups = aggregate_summary(GroupBy::Source, &records, None);
        let diff = groups
            .iter()
            .find(|group| group.key == "diff")
            .expect("diff");
        assert_eq!(diff.tokens.input, 340);
        assert_eq!(diff.sessions, 1);
        assert_eq!(diff.turns, 0);
    }

    #[test]
    fn aggregate_timeseries_groups_sessions_by_day() {
        let created_a = parse_rfc3339("2026-02-01T09:00:00Z");
        let created_b = parse_rfc3339("2026-02-02T11:00:00Z");
        let records = vec![
            UsageSessionData {
                session: test_session_at("implement", created_a),
                events: vec![turn_usage_event(TurnUsage {
                    input_tokens: 100,
                    output_tokens: 20,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                })],
                wave_id: Some("wave-alpha".to_string()),
                flow: Some("build".to_string()),
            },
            UsageSessionData {
                session: test_session_at("implement", created_b),
                events: vec![turn_usage_event(TurnUsage {
                    input_tokens: 40,
                    output_tokens: 10,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                })],
                wave_id: Some("wave-alpha".to_string()),
                flow: Some("build".to_string()),
            },
        ];

        let buckets = aggregate_timeseries(TimeBucket::Day, GroupBy::Wave, records, None);
        assert_eq!(buckets.len(), 2);
        assert_eq!(buckets[0].period, "2026-02-01");
        assert_eq!(buckets[1].period, "2026-02-02");
        assert_eq!(buckets[0].groups[0].key, "wave-alpha");
        assert_eq!(buckets[0].groups[0].tokens.input, 100);
        assert_eq!(buckets[1].groups[0].tokens.input, 40);
    }

    #[test]
    fn aggregate_timeseries_week_bucket_uses_monday_start() {
        let monday = parse_rfc3339("2026-02-02T08:00:00Z");
        let wednesday = parse_rfc3339("2026-02-04T12:00:00Z");
        let records = vec![
            UsageSessionData {
                session: test_session_at("implement", monday),
                events: vec![turn_usage_event(TurnUsage {
                    input_tokens: 60,
                    output_tokens: 10,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                })],
                wave_id: Some("wave-alpha".to_string()),
                flow: Some("build".to_string()),
            },
            UsageSessionData {
                session: test_session_at("implement", wednesday),
                events: vec![turn_usage_event(TurnUsage {
                    input_tokens: 90,
                    output_tokens: 12,
                    reasoning_tokens: None,
                    cache_read_tokens: None,
                    cache_write_tokens: None,
                    model: Some("claude-sonnet-4".to_string()),
                    cost_usd: None,
                })],
                wave_id: Some("wave-alpha".to_string()),
                flow: Some("build".to_string()),
            },
        ];

        let buckets = aggregate_timeseries(TimeBucket::Week, GroupBy::Wave, records, None);
        assert_eq!(buckets.len(), 1);
        assert_eq!(buckets[0].period, "2026-02-02");
        assert_eq!(buckets[0].groups[0].tokens.input, 150);
        assert_eq!(buckets[0].groups[0].turns, 2);
    }

    fn test_session(step: &str) -> Session {
        test_session_at(step, OffsetDateTime::now_utc())
    }

    fn test_session_at(step: &str, created_at: OffsetDateTime) -> Session {
        Session {
            id: LfdId::new(),
            harness: "claude".to_string(),
            status: SessionStatus::Ended,
            wave_run_id: None,
            provider_session_id: None,
            config: SessionConfig {
                step: step.to_string(),
                repo_root: "/tmp/repo".to_string(),
                ..Default::default()
            },
            created_at,
            ended_at: None,
        }
    }

    fn parse_rfc3339(value: &str) -> OffsetDateTime {
        OffsetDateTime::parse(value, &Rfc3339).expect("parse timestamp")
    }

    fn context_event(sources: HashMap<String, u64>) -> PersistedSessionEvent {
        PersistedSessionEvent {
            session_id: LfdId::new(),
            seq: 0,
            event: SessionEvent::ContextSnapshot {
                snapshot: ContextSnapshot {
                    sources,
                    source_counts: HashMap::new(),
                    documents: Vec::new(),
                    budget: 200_000,
                    total: 460,
                    diff_tier: "UnifiedDiff".to_string(),
                    step_name: None,
                    direction_names: Vec::new(),
                    area_name: None,
                    wave_name: None,
                    has_clipboard: false,
                },
            },
            created_at: OffsetDateTime::now_utc(),
        }
    }

    fn turn_usage_event(usage: TurnUsage) -> PersistedSessionEvent {
        PersistedSessionEvent {
            session_id: LfdId::new(),
            seq: 1,
            event: SessionEvent::TurnUsage {
                turn_id: "turn_1".to_string(),
                usage,
            },
            created_at: OffsetDateTime::now_utc(),
        }
    }
}
