jobtrackingapp/JobTrackerApi/Services/AiUsageMeter.cs

using JobTrackerApi.Data;
using JobTrackerApi.Models;
using Microsoft.EntityFrameworkCore;

namespace JobTrackerApi.Services;

public sealed record AiUsageTotals(int Calls, long InputCharacters, long OutputCharacters, long EstimatedTokens);
public sealed record AiUsageReservation(AiUsageRecord Record, bool Created);

public sealed class AiUsageLimitException(string code, string message) : Exception(message)
{
    public string Code { get; } = code;
}

public sealed class AiUsageMeter(JobTrackerContext db, TimeProvider timeProvider)
{
    private static readonly SemaphoreSlim Gate = new(1, 1);

    public static (int InputCharacters, int EstimatedTokens) ReservationFor(string taskType) => taskType switch
    {
        StrategySnapshotService.TaskType => (48_000, 12_000),
        CvProcessingQueue.TaskType => (64_000, 16_000),
        _ => (16_000, 4_000),
    };

    public async Task<AiUsageTotals> CurrentMonthAsync(string ownerUserId, CancellationToken cancellationToken)
        => await SinceAsync(ownerUserId, MonthStart(timeProvider.GetUtcNow()), cancellationToken);

    public async Task<AiUsageTotals> AllTimeAsync(string ownerUserId, CancellationToken cancellationToken)
        => await SinceAsync(ownerUserId, null, cancellationToken);

    public async Task EnsureCanReserveAsync(
        string ownerUserId,
        AccountEntitlements entitlements,
        int calls,
        int estimatedTokens,
        CancellationToken cancellationToken)
    {
        var used = await CurrentMonthAsync(ownerUserId, cancellationToken);
        EnsureWithinLimit(used, entitlements, calls, estimatedTokens);
    }

    public async Task<AiUsageReservation> ReserveAsync(
        string ownerUserId,
        AccountEntitlements entitlements,
        string sourceType,
        string sourceId,
        string taskType,
        int inputCharacters,
        int estimatedTokens,
        CancellationToken cancellationToken)
    {
        Validate(sourceType, sourceId, taskType, inputCharacters, estimatedTokens);
        await Gate.WaitAsync(cancellationToken);
        try
        {
            var existing = await db.AiUsageRecords.FirstOrDefaultAsync(
                item => item.SourceType == sourceType && item.SourceId == sourceId,
                cancellationToken);
            if (existing is not null) return new AiUsageReservation(existing, false);

            await EnsureCanReserveAsync(ownerUserId, entitlements, 1, estimatedTokens, cancellationToken);
            var record = NewRecord(ownerUserId, sourceType, sourceId, taskType, inputCharacters, estimatedTokens, timeProvider.GetUtcNow());
            db.AiUsageRecords.Add(record);
            await db.SaveChangesAsync(cancellationToken);
            return new AiUsageReservation(record, true);
        }
        finally
        {
            Gate.Release();
        }
    }

    public async Task FinalizeAsync(long id, int inputCharacters, int outputCharacters, CancellationToken cancellationToken)
    {
        if (inputCharacters < 0 || outputCharacters < 0) throw new ArgumentOutOfRangeException();
        var estimatedTokens = (inputCharacters + outputCharacters + 3) / 4;
        await db.AiUsageRecords.Where(item => item.Id == id).ExecuteUpdateAsync(setters => setters
            .SetProperty(item => item.InputCharacterCount, inputCharacters)
            .SetProperty(item => item.OutputCharacterCount, outputCharacters)
            .SetProperty(item => item.EstimatedTokenCount, estimatedTokens), cancellationToken);
    }

    public async Task ReleaseAsync(long id, CancellationToken cancellationToken)
        => await db.AiUsageRecords.Where(item => item.Id == id).ExecuteDeleteAsync(cancellationToken);

    public static AiUsageRecord NewOperationRecord(
        string ownerUserId,
        Guid operationId,
        string taskType,
        int inputCharacters,
        int estimatedTokens,
        DateTimeOffset createdAtUtc)
        => NewRecord(ownerUserId, "operation", operationId.ToString("D"), taskType, inputCharacters, estimatedTokens, createdAtUtc);

    private async Task<AiUsageTotals> SinceAsync(string ownerUserId, DateTimeOffset? since, CancellationToken cancellationToken)
    {
        var query = db.AiUsageRecords.Where(item => item.OwnerUserId == ownerUserId);
        if (db.Database.IsSqlite())
        {
            var rows = await query.AsNoTracking().ToListAsync(cancellationToken);
            if (since is not null) rows = rows.Where(item => item.CreatedAtUtc >= since.Value).ToList();
            return Sum(rows);
        }

        if (since is not null) query = query.Where(item => item.CreatedAtUtc >= since.Value);
        var totals = await query.GroupBy(_ => 1).Select(group => new AiUsageTotals(
            group.Sum(item => item.CallCount),
            group.Sum(item => (long)item.InputCharacterCount),
            group.Sum(item => (long)item.OutputCharacterCount),
            group.Sum(item => (long)item.EstimatedTokenCount))).FirstOrDefaultAsync(cancellationToken);
        return totals ?? new AiUsageTotals(0, 0, 0, 0);
    }

    private static AiUsageTotals Sum(IEnumerable<AiUsageRecord> records) => new(
        records.Sum(item => item.CallCount),
        records.Sum(item => (long)item.InputCharacterCount),
        records.Sum(item => (long)item.OutputCharacterCount),
        records.Sum(item => (long)item.EstimatedTokenCount));

    private static void EnsureWithinLimit(AiUsageTotals used, AccountEntitlements entitlements, int calls, int tokens)
    {
        if (used.Calls + calls > entitlements.MonthlyAiCalls)
            throw new AiUsageLimitException("monthly_ai_calls_exhausted", $"Monthly AI limit reached ({entitlements.MonthlyAiCalls} generations). Try again next month.");
        if (used.EstimatedTokens + tokens > entitlements.MonthlyAiTokens)
            throw new AiUsageLimitException("monthly_ai_tokens_exhausted", $"Monthly AI cost limit reached ({entitlements.MonthlyAiTokens:N0} estimated tokens). Try again next month.");
    }

    private static AiUsageRecord NewRecord(string ownerUserId, string sourceType, string sourceId, string taskType, int inputCharacters, int estimatedTokens, DateTimeOffset createdAtUtc)
    {
        Validate(sourceType, sourceId, taskType, inputCharacters, estimatedTokens);
        return new AiUsageRecord
        {
            OwnerUserId = ownerUserId,
            SourceType = sourceType,
            SourceId = sourceId,
            TaskType = taskType,
            InputCharacterCount = inputCharacters,
            EstimatedTokenCount = estimatedTokens,
            CreatedAtUtc = createdAtUtc,
        };
    }

    private static void Validate(string sourceType, string sourceId, string taskType, int inputCharacters, int estimatedTokens)
    {
        if (string.IsNullOrWhiteSpace(sourceType) || sourceType.Length > 32) throw new ArgumentOutOfRangeException(nameof(sourceType));
        if (string.IsNullOrWhiteSpace(sourceId) || sourceId.Length > 64) throw new ArgumentOutOfRangeException(nameof(sourceId));
        if (string.IsNullOrWhiteSpace(taskType) || taskType.Length > 64) throw new ArgumentOutOfRangeException(nameof(taskType));
        if (inputCharacters < 0) throw new ArgumentOutOfRangeException(nameof(inputCharacters));
        if (estimatedTokens < 0) throw new ArgumentOutOfRangeException(nameof(estimatedTokens));
    }

    private static DateTimeOffset MonthStart(DateTimeOffset value)
        => new(value.Year, value.Month, 1, 0, 0, 0, TimeSpan.Zero);
}