KArtSell.Aegis/src/KArtSell.Host/Infrastructure/RateLimiterService.cs

using Dapper;
using KArtSell.BuildingBlocks.Time;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Logging;
using Npgsql;

namespace KArtSell.Host.Infrastructure;

/// <summary>
/// Central rate limiter using token bucket pattern.
/// Per-API quotas: KRX 100/min, OpenDart 1000/day, KIS 50/sec.
/// Atomic token consumption, no partial success, HTTP 429 with retry-after header.
/// </summary>
public class RateLimiterService
{
    private readonly NpgsqlDataSource _dataSource;
    private readonly IClock _clock;
    private readonly ILogger<RateLimiterService> _logger;

    private static readonly Dictionary<string, RateLimitConfig> ApiConfigs = new()
    {
        { "krx", new RateLimitConfig { Limit = 100, WindowSeconds = 60 } },
        { "opendart", new RateLimitConfig { Limit = 1000, WindowSeconds = 86400 } },  // 1 day
        { "kis", new RateLimitConfig { Limit = 50, WindowSeconds = 1 } }
    };

    private static readonly Action<ILogger, string, int, Exception?> LogTokenConsumed =
        LoggerMessage.Define<string, int>(
            LogLevel.Debug,
            new EventId(1, nameof(LogTokenConsumed)),
            "Rate limit: {ApiName} consumed 1 token, {RemainTokens} remaining");

    private static readonly Action<ILogger, string, int, Exception?> LogQuotaExceeded =
        LoggerMessage.Define<string, int>(
            LogLevel.Warning,
            new EventId(2, nameof(LogQuotaExceeded)),
            "Rate limit: {ApiName} quota exceeded, retry after {RetryAfter}s");

    public RateLimiterService(NpgsqlDataSource dataSource, IClock clock, ILogger<RateLimiterService> logger)
    {
        _dataSource = dataSource;
        _clock = clock;
        _logger = logger;
    }

    /// <summary>
    /// Attempt to consume 1 token from the rate limit bucket for the given API.
    /// Returns true if successful; false if quota exhausted.
    /// </summary>
    public async Task<(bool Success, int RetryAfterSeconds)> TryConsumeAsync(
        string apiName,
        CancellationToken cancellationToken = default)
    {
        if (!ApiConfigs.TryGetValue(apiName.ToLowerInvariant(), out var config))
        {
            _logger.LogWarning("Unknown API for rate limiting: {ApiName}", apiName);
            return (false, 0);
        }

        // Atomic consumption in database
        const string sql = """
            UPDATE infrastructure.rate_limit_quota
            SET current_tokens = current_tokens - 1, updated_at = @now
            WHERE api_name = @apiName
              AND current_tokens > 0
            RETURNING current_tokens, window_seconds, limit_count
            """;

        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken);
        var result = await connection.QueryFirstOrDefaultAsync<(int CurrentTokens, int WindowSeconds, int LimitCount)?>(
            sql,
            new { apiName = apiName.ToLower(), now = _clock.UtcNow.UtcDateTime },
            commandTimeout: 5);

        if (result == null)
        {
            // Quota exhausted
            var retryAfter = config.WindowSeconds;
            LogQuotaExceeded(_logger, apiName, retryAfter, null);

            // Log rejection event
            await LogEventAsync(apiName, "rejected", 1, 0, 0, cancellationToken);

            return (false, retryAfter);
        }

        // Token consumed successfully
        LogTokenConsumed(_logger, apiName, result.Value.CurrentTokens, null);
        await LogEventAsync(apiName, "allowed", 1, 1, result.Value.CurrentTokens, cancellationToken);

        return (true, 0);
    }

    /// <summary>
    /// Reset quota for the given API (e.g., daily reset for OpenDart).
    /// Called by scheduled job at window boundary.
    /// </summary>
    public async Task ResetQuotaAsync(string apiName, CancellationToken cancellationToken = default)
    {
        if (!ApiConfigs.TryGetValue(apiName.ToLowerInvariant(), out var config))
            return;

        const string sql = """
            UPDATE infrastructure.rate_limit_quota
            SET current_tokens = @limit, last_reset_at = @now, updated_at = @now
            WHERE api_name = @apiName
            """;

        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken);
        await connection.ExecuteAsync(
            sql,
            new { apiName = apiName.ToLower(), limit = config.Limit, now = _clock.UtcNow.UtcDateTime },
            commandTimeout: 5);

        _logger.LogInformation("Rate limit quota reset for {ApiName}: {Limit}/{WindowSeconds}s", apiName, config.Limit, config.WindowSeconds);
    }

    /// <summary>
    /// Initialize rate limit quotas (called at startup).
    /// </summary>
    public async Task InitializeAsync(CancellationToken cancellationToken = default)
    {
        const string sql = """
            INSERT INTO infrastructure.rate_limit_quota (api_name, limit_count, window_seconds, current_tokens, last_reset_at, updated_at, published_at)
            VALUES (@apiName, @limit, @window, @limit, @now, @now, @now)
            ON CONFLICT (api_name) DO NOTHING
            """;

        await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken);

        foreach (var (apiName, config) in ApiConfigs)
        {
            await connection.ExecuteAsync(
                sql,
                new { apiName, limit = config.Limit, window = config.WindowSeconds, now = _clock.UtcNow.UtcDateTime },
                commandTimeout: 5);
        }

        _logger.LogInformation("Rate limit quotas initialized: {Count} APIs", ApiConfigs.Count);
    }

    private async Task LogEventAsync(string apiName, string decision, int tokensRequested, int tokensUsed, decimal remainingTokens, CancellationToken cancellationToken)
    {
        const string sql = """
            INSERT INTO infrastructure.rate_limit_events (api_name, decision, tokens_requested, tokens_used, remaining_tokens, occurred_at, published_at)
            VALUES (@apiName, @decision, @tokensRequested, @tokensUsed, @remainingTokens, @now, @now)
            """;

        try
        {
            await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken);
            await connection.ExecuteAsync(
                sql,
                new { apiName = apiName.ToLower(), decision, tokensRequested, tokensUsed, remainingTokens, now = _clock.UtcNow.UtcDateTime },
                commandTimeout: 5);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to log rate limit event for {ApiName}", apiName);
        }
    }
}

public class RateLimitConfig
{
    public int Limit { get; set; }
    public int WindowSeconds { get; set; }
}

/// <summary>
/// Middleware: Apply rate limiting to incoming HTTP requests.
/// Returns HTTP 429 (Too Many Requests) if quota exhausted.
/// </summary>
public class RateLimiterMiddleware
{
    private readonly RequestDelegate _next;
    private readonly ILogger<RateLimiterMiddleware> _logger;

    public RateLimiterMiddleware(RequestDelegate next, ILogger<RateLimiterMiddleware> logger)
    {
        _next = next;
        _logger = logger;
    }

    public async Task InvokeAsync(HttpContext context, RateLimiterService rateLimiter)
    {
        // Determine API from route (e.g., /api/krx/* → krx)
        var path = context.Request.Path.Value?.ToLower() ?? "";
        string? apiName = null;

        if (path.Contains("/krx/")) apiName = "krx";
        else if (path.Contains("/opendart/")) apiName = "opendart";
        else if (path.Contains("/kis/")) apiName = "kis";

        // Only rate limit if API is identified
        if (apiName != null)
        {
            var (success, retryAfter) = await rateLimiter.TryConsumeAsync(apiName, context.RequestAborted);

            if (!success)
            {
                context.Response.StatusCode = StatusCodes.Status429TooManyRequests;
                context.Response.Headers.Add("Retry-After", retryAfter.ToString());
                await context.Response.WriteAsync($"Rate limit exceeded for {apiName}. Retry after {retryAfter}s");
                return;
            }
        }

        await _next(context);
    }
}