namespace ClaudeDo.Worker.Usage.TokenTracker;

public sealed record TokenTrackerTotals(long TotalTokens, double CostUsd, int Sessions);

/// <summary>
/// Pure aggregation over a parsed export. One fetch covers the whole history — <c>--from</c> and
/// <c>--to</c> are accepted by the CLI but do <b>not</b> filter its output (measured against
/// v0.88.4: identical 886 rows with and without a range) — so filtering here is load-bearing, and
/// switching the range in the UI costs nothing.
/// </summary>
public static class TokenTrackerAggregator
{
    public const string ScopeClaudeDo = "claudedo";
    public const string ScopeOther = "other";

    public static IReadOnlyList<TokenTrackerModelRow> ByModel(
        IReadOnlyList<TokenTrackerSession> sessions,
        DateOnly from,
        DateOnly to,
        ISet<string> claudeDoHashes)
    {
        var buckets = new Dictionary<(DateOnly Date, string Model, string Scope), Accumulator>();

        foreach (var session in InRange(sessions, from, to))
        {
            var key = (
                DateOf(session),
                session.Model,
                claudeDoHashes.Contains(session.SessionHash) ? ScopeClaudeDo : ScopeOther);

            if (!buckets.TryGetValue(key, out var acc))
            {
                acc = new Accumulator();
                buckets[key] = acc;
            }

            acc.Input += session.InputTokens;
            acc.Output += session.OutputTokens;
            acc.CacheRead += session.CacheReadTokens;
            acc.CacheCreation += session.CacheCreationTokens;
            acc.Cost += session.CostUsd;
            acc.Sessions++;
        }

        return buckets
            .Select(kv => new TokenTrackerModelRow(
                kv.Key.Date, kv.Key.Model, kv.Key.Scope,
                kv.Value.Input, kv.Value.Output, kv.Value.CacheRead, kv.Value.CacheCreation,
                kv.Value.Sessions, kv.Value.Cost))
            .OrderBy(r => r.Date).ThenBy(r => r.Model).ThenBy(r => r.Scope)
            .ToList();
    }

    /// <summary>Cost and efficiency for one task, given the session ids of its runs. Null when
    /// none of them appear in the export — the caller then shows tokens without cost. Several
    /// export rows can share one session hash (they are disjoint segments of the same session),
    /// so summing every match is correct rather than double-counting.</summary>
    public static TokenTrackerTaskExtras? ForSessions(
        IReadOnlyList<TokenTrackerSession> sessions,
        IEnumerable<string?> sessionIds)
    {
        var wanted = sessionIds
            .Select(SessionHash.ForClaudeSession)
            .Where(h => h is not null)
            .ToHashSet()!;

        if (wanted.Count == 0) return null;

        var matched = sessions.Where(s => wanted.Contains(s.SessionHash)).ToList();
        if (matched.Count == 0) return null;

        return new TokenTrackerTaskExtras(
            CostUsd: matched.Sum(s => s.CostUsd),
            Retries: matched.Sum(s => s.RetryTurns),
            Productive: matched.Any(s => s.Productive),
            OneShot: matched.All(s => s.OneShot));
    }

    public static TokenTrackerTotals Totals(
        IReadOnlyList<TokenTrackerSession> sessions, DateOnly from, DateOnly to)
    {
        var inRange = InRange(sessions, from, to).ToList();
        return new TokenTrackerTotals(
            TotalTokens: inRange.Sum(s =>
                s.InputTokens + s.OutputTokens + s.CacheReadTokens + s.CacheCreationTokens),
            CostUsd: inRange.Sum(s => s.CostUsd),
            Sessions: inRange.Count);
    }

    private static IEnumerable<TokenTrackerSession> InRange(
        IReadOnlyList<TokenTrackerSession> sessions, DateOnly from, DateOnly to) =>
        sessions.Where(s => DateOf(s) >= from && DateOf(s) <= to);

    /// <summary>Local start date — the UI's pickers and the old transcript reader both work in
    /// local time, so a UTC-based date would shift rows across midnight for the user.</summary>
    private static DateOnly DateOf(TokenTrackerSession session) =>
        DateOnly.FromDateTime(session.StartedAt.LocalDateTime);

    private sealed class Accumulator
    {
        public long Input, Output, CacheRead, CacheCreation;
        public double Cost;
        public int Sessions;
    }
}
