gitcommiteditor/Services/CostTrackingService.cs

using System.Collections.Concurrent;
using System.Text.Json;

namespace MarketAlly.GitCommitEditor.Services;

/// <summary>
/// Persistence provider interface for cost tracking
/// </summary>
public interface ICostPersistenceProvider
{
    string? GetValue(string key);
    void SetValue(string key, string value);
    void Remove(string key);
}

/// <summary>
/// Tracks AI operation costs for the session with persistence support
/// </summary>
public class CostTrackingService : ICostTrackingService
{
    private const string KeyLifetimeCost = "cost_lifetime_total";
    private const string KeyLifetimeOperations = "cost_lifetime_operations";
    private const string KeyLifetimeInputTokens = "cost_lifetime_input_tokens";
    private const string KeyLifetimeOutputTokens = "cost_lifetime_output_tokens";

    private readonly ConcurrentDictionary<string, OperationCostSummary> _costBreakdown = new();
    private readonly ICostPersistenceProvider? _persistence;
    private readonly object _costLock = new();

    private decimal _sessionCost;
    private int _totalInputTokens;
    private int _totalOutputTokens;
    private int _operationCount;

    private decimal _lifetimeCost;
    private int _lifetimeInputTokens;
    private int _lifetimeOutputTokens;
    private int _lifetimeOperationCount;

    public decimal SessionCost => _sessionCost;
    public decimal LifetimeCost => _lifetimeCost;
    public int TotalInputTokens => _totalInputTokens;
    public int TotalOutputTokens => _totalOutputTokens;
    public int OperationCount => _operationCount;
    public int LifetimeOperationCount => _lifetimeOperationCount;

    public event EventHandler<CostUpdatedEventArgs>? CostUpdated;

    public CostTrackingService(ICostPersistenceProvider? persistence = null)
    {
        _persistence = persistence;
        LoadState();
    }

    public void RecordOperation(string operationType, int inputTokens, int outputTokens, decimal cost)
    {
        // Update session totals
        Interlocked.Add(ref _totalInputTokens, inputTokens);
        Interlocked.Add(ref _totalOutputTokens, outputTokens);
        Interlocked.Increment(ref _operationCount);

        // Update lifetime totals
        Interlocked.Add(ref _lifetimeInputTokens, inputTokens);
        Interlocked.Add(ref _lifetimeOutputTokens, outputTokens);
        Interlocked.Increment(ref _lifetimeOperationCount);

        // Thread-safe decimal addition using lock
        lock (_costLock)
        {
            _sessionCost += cost;
            _lifetimeCost += cost;
        }

        // Update breakdown
        _costBreakdown.AddOrUpdate(
            operationType,
            _ => new OperationCostSummary
            {
                OperationType = operationType,
                Count = 1,
                TotalInputTokens = inputTokens,
                TotalOutputTokens = outputTokens,
                TotalCost = cost
            },
            (_, existing) =>
            {
                existing.Count++;
                existing.TotalInputTokens += inputTokens;
                existing.TotalOutputTokens += outputTokens;
                existing.TotalCost += cost;
                return existing;
            });

        // Auto-save after each operation
        SaveState();

        // Raise event
        CostUpdated?.Invoke(this, new CostUpdatedEventArgs(
            operationType,
            cost,
            _sessionCost,
            inputTokens,
            outputTokens));
    }

    public IReadOnlyDictionary<string, OperationCostSummary> GetCostBreakdown()
    {
        return _costBreakdown.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
    }

    public void ResetSession()
    {
        _costBreakdown.Clear();
        _sessionCost = 0;
        _totalInputTokens = 0;
        _totalOutputTokens = 0;
        _operationCount = 0;

        CostUpdated?.Invoke(this, new CostUpdatedEventArgs("SessionReset", 0, 0, 0, 0));
    }

    public void ResetAll()
    {
        ResetSession();
        _lifetimeCost = 0;
        _lifetimeInputTokens = 0;
        _lifetimeOutputTokens = 0;
        _lifetimeOperationCount = 0;

        // Clear persisted data
        _persistence?.Remove(KeyLifetimeCost);
        _persistence?.Remove(KeyLifetimeOperations);
        _persistence?.Remove(KeyLifetimeInputTokens);
        _persistence?.Remove(KeyLifetimeOutputTokens);

        CostUpdated?.Invoke(this, new CostUpdatedEventArgs("AllReset", 0, 0, 0, 0));
    }

    public void SaveState()
    {
        if (_persistence == null) return;

        try
        {
            _persistence.SetValue(KeyLifetimeCost, _lifetimeCost.ToString("G"));
            _persistence.SetValue(KeyLifetimeOperations, _lifetimeOperationCount.ToString());
            _persistence.SetValue(KeyLifetimeInputTokens, _lifetimeInputTokens.ToString());
            _persistence.SetValue(KeyLifetimeOutputTokens, _lifetimeOutputTokens.ToString());
        }
        catch
        {
            // Silently fail - persistence is optional
        }
    }

    public void LoadState()
    {
        if (_persistence == null) return;

        try
        {
            var costStr = _persistence.GetValue(KeyLifetimeCost);
            if (!string.IsNullOrEmpty(costStr) && decimal.TryParse(costStr, out var cost))
            {
                _lifetimeCost = cost;
            }

            var opsStr = _persistence.GetValue(KeyLifetimeOperations);
            if (!string.IsNullOrEmpty(opsStr) && int.TryParse(opsStr, out var ops))
            {
                _lifetimeOperationCount = ops;
            }

            var inputStr = _persistence.GetValue(KeyLifetimeInputTokens);
            if (!string.IsNullOrEmpty(inputStr) && int.TryParse(inputStr, out var input))
            {
                _lifetimeInputTokens = input;
            }

            var outputStr = _persistence.GetValue(KeyLifetimeOutputTokens);
            if (!string.IsNullOrEmpty(outputStr) && int.TryParse(outputStr, out var output))
            {
                _lifetimeOutputTokens = output;
            }
        }
        catch
        {
            // Silently fail - start fresh if persistence fails
        }
    }
}