ClaudeDo/src/ClaudeDo.Worker/Usage/TranscriptUsageReader.cs

using System.Collections.Concurrent;
using System.Text.Json;
using ClaudeDo.Data;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Usage.Interfaces;

namespace ClaudeDo.Worker.Usage;

public sealed class TranscriptUsageReader : ITranscriptUsageReader
{
    private readonly string _projectsRoot;
    private readonly string _centralRoot;
    private readonly string _sandboxRoot;
    private readonly ConcurrentDictionary<string, FileCacheEntry> _cache = new();

    public TranscriptUsageReader(WorkerConfig cfg, string? projectsRoot = null)
    {
        _projectsRoot = projectsRoot ?? Paths.Expand("~/.claude/projects");
        _centralRoot = NormalizePath(cfg.CentralWorktreeRoot);
        _sandboxRoot = NormalizePath(cfg.SandboxRoot);
    }

    public Task<IReadOnlyList<UsageAggregateRow>> ReadAsync(
        DateOnly start, DateOnly end, CancellationToken ct = default)
    {
        var seenKeys = new HashSet<string>();
        var buckets = new Dictionary<(DateOnly Date, string Model, UsageScope Scope), Accumulator>();

        if (Directory.Exists(_projectsRoot))
        {
            foreach (var file in Directory.EnumerateFiles(_projectsRoot, "*.jsonl", SearchOption.AllDirectories))
            {
                ct.ThrowIfCancellationRequested();

                foreach (var record in GetOrReadFile(file))
                {
                    if (record.Date < start || record.Date > end) continue;
                    if (!seenKeys.Add(record.DedupeKey)) continue;

                    var key = (record.Date, record.Model, record.Scope);
                    if (!buckets.TryGetValue(key, out var acc))
                    {
                        acc = new Accumulator();
                        buckets[key] = acc;
                    }
                    acc.Input += record.InputTokens;
                    acc.Output += record.OutputTokens;
                    acc.CacheRead += record.CacheReadTokens;
                    acc.CacheCreation += record.CacheCreationTokens;
                    acc.Messages++;
                }
            }
        }

        var rows = buckets
            .Select(kv => new UsageAggregateRow(
                kv.Key.Date, kv.Key.Model, kv.Key.Scope,
                kv.Value.Input, kv.Value.Output, kv.Value.CacheRead, kv.Value.CacheCreation, kv.Value.Messages))
            .OrderBy(r => r.Date).ThenBy(r => r.Model).ThenBy(r => r.Scope)
            .ToList();

        return Task.FromResult<IReadOnlyList<UsageAggregateRow>>(rows);
    }

    public Task<SessionUsageTotals?> ReadSessionTotalsAsync(string sessionId, CancellationToken ct = default)
    {
        if (string.IsNullOrWhiteSpace(sessionId) || !Directory.Exists(_projectsRoot))
            return Task.FromResult<SessionUsageTotals?>(null);

        var file = Directory
            .EnumerateFiles(_projectsRoot, $"{sessionId}.jsonl", SearchOption.AllDirectories)
            .FirstOrDefault();
        if (file is null) return Task.FromResult<SessionUsageTotals?>(null);

        var seenKeys = new HashSet<string>();
        long input = 0, output = 0, cacheRead = 0, cacheCreation = 0;
        foreach (var record in GetOrReadFile(file))
        {
            ct.ThrowIfCancellationRequested();
            if (!seenKeys.Add(record.DedupeKey)) continue;

            input += record.InputTokens;
            output += record.OutputTokens;
            cacheRead += record.CacheReadTokens;
            cacheCreation += record.CacheCreationTokens;
        }

        return Task.FromResult<SessionUsageTotals?>(
            new SessionUsageTotals(input, output, cacheRead, cacheCreation));
    }

    private List<UsageMessageRecord> GetOrReadFile(string file)
    {
        var info = new FileInfo(file);
        if (_cache.TryGetValue(file, out var cached) &&
            cached.Length == info.Length && cached.LastWriteUtc == info.LastWriteTimeUtc)
        {
            return cached.Records;
        }

        var records = ReadFile(file);
        _cache[file] = new FileCacheEntry(info.Length, info.LastWriteTimeUtc, records);
        return records;
    }

    private List<UsageMessageRecord> ReadFile(string file)
    {
        var records = new List<UsageMessageRecord>();

        foreach (var line in File.ReadLines(file))
        {
            if (string.IsNullOrWhiteSpace(line)) continue;

            JsonDocument doc;
            try { doc = JsonDocument.Parse(line); }
            catch (JsonException) { continue; }

            using (doc)
            {
                var root = doc.RootElement;
                if (root.ValueKind != JsonValueKind.Object) continue;
                if (!root.TryGetProperty("type", out var typeEl) || typeEl.GetString() != "assistant") continue;
                if (!root.TryGetProperty("timestamp", out var tsEl) ||
                    !DateTimeOffset.TryParse(tsEl.GetString(), out var ts)) continue;
                if (!root.TryGetProperty("message", out var msg) || msg.ValueKind != JsonValueKind.Object) continue;
                if (!msg.TryGetProperty("model", out var modelEl) || modelEl.ValueKind != JsonValueKind.String) continue;

                var date = DateOnly.FromDateTime(ts.LocalDateTime);
                var model = modelEl.GetString()!;
                if (model == "<synthetic>") continue;

                long input = 0, output = 0, cacheRead = 0, cacheCreation = 0;
                if (msg.TryGetProperty("usage", out var usage) && usage.ValueKind == JsonValueKind.Object)
                {
                    input = GetLong(usage, "input_tokens");
                    output = GetLong(usage, "output_tokens");
                    cacheRead = GetLong(usage, "cache_read_input_tokens");
                    cacheCreation = GetLong(usage, "cache_creation_input_tokens");
                }

                var cwd = TryGetString(root, "cwd") ?? "";
                var dedupeKey = TryGetString(root, "requestId")
                    ?? TryGetString(msg, "id")
                    ?? Guid.NewGuid().ToString();

                records.Add(new UsageMessageRecord(
                    date, model, ResolveScope(cwd), input, output, cacheRead, cacheCreation, dedupeKey));
            }
        }

        return records;
    }

    private UsageScope ResolveScope(string cwd)
    {
        var norm = NormalizePath(cwd);
        if (norm.Length == 0) return UsageScope.Other;

        // Sibling-strategy worktrees are placed next to whatever repo they belong to
        // (no single root path), but always under a literal ".claudedo-worktrees" segment.
        if (norm.Split('\\').Any(seg => seg == ".claudedo-worktrees")) return UsageScope.ClaudeDo;
        if (IsUnderRoot(norm, _centralRoot)) return UsageScope.ClaudeDo;
        if (IsUnderRoot(norm, _sandboxRoot)) return UsageScope.ClaudeDo;
        return UsageScope.Other;
    }

    private static bool IsUnderRoot(string normPath, string normRoot) =>
        normRoot.Length > 0 && (normPath == normRoot || normPath.StartsWith(normRoot + "\\", StringComparison.Ordinal));

    private static string? TryGetString(JsonElement obj, string prop) =>
        obj.TryGetProperty(prop, out var el) && el.ValueKind == JsonValueKind.String ? el.GetString() : null;

    private static long GetLong(JsonElement obj, string prop) =>
        obj.TryGetProperty(prop, out var el) && el.TryGetInt64(out var v) ? v : 0;

    private static string NormalizePath(string p) =>
        (p ?? "").Replace('/', '\\').TrimEnd('\\').ToLowerInvariant();

    private sealed class Accumulator
    {
        public long Input, Output, CacheRead, CacheCreation;
        public int Messages;
    }

    private sealed record FileCacheEntry(long Length, DateTime LastWriteUtc, List<UsageMessageRecord> Records);

    private sealed record UsageMessageRecord(
        DateOnly Date, string Model, UsageScope Scope,
        long InputTokens, long OutputTokens, long CacheReadTokens, long CacheCreationTokens,
        string DedupeKey);
}