using System.Collections.Concurrent; using System.Text.Json; namespace MarketAlly.GitCommitEditor.Services; /// /// Persistence provider interface for cost tracking /// public interface ICostPersistenceProvider { string? GetValue(string key); void SetValue(string key, string value); void Remove(string key); } /// /// Tracks AI operation costs for the session with persistence support /// 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 _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? 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 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 } } }