using System.Threading.RateLimiting;
using Microsoft.AspNetCore.Mvc;
using Microsoft.AspNetCore.RateLimiting;
using Tiku.Api.Middleware;
using Tiku.Api.Options;
using Tiku.Application.Auth;
using Tiku.Application.Security;

namespace Tiku.Api.Configuration;

internal static class RateLimitingExtensions
{
    internal static IServiceCollection AddApiRateLimiting(
        this IServiceCollection services,
        IConfiguration configuration)
    {
        services.AddOptions<ApiRateLimitOptions>()
            .Bind(configuration.GetSection(ApiRateLimitOptions.SectionName))
            .ValidateDataAnnotations()
            .ValidateOnStart();
        var rateLimitOptions = configuration
            .GetSection(ApiRateLimitOptions.SectionName)
            .Get<ApiRateLimitOptions>() ?? new ApiRateLimitOptions();

        services.AddOptions<AuthRateLimitOptions>()
            .Bind(configuration.GetSection(AuthRateLimitOptions.SectionName))
            .ValidateDataAnnotations()
            .ValidateOnStart();
        var authRateLimitOptions = configuration
            .GetSection(AuthRateLimitOptions.SectionName)
            .Get<AuthRateLimitOptions>() ?? new AuthRateLimitOptions();

        services.AddRateLimiter(options =>
        {
            options.RejectionStatusCode = StatusCodes.Status429TooManyRequests;
            if (rateLimitOptions.Enabled)
            {
                options.GlobalLimiter = PartitionedRateLimiter.Create<HttpContext, string>(httpContext =>
                {
                    var partitionKey =
                        httpContext.User.FindFirst(TikuClaimTypes.UserId)?.Value ??
                        httpContext.Connection.RemoteIpAddress?.ToString() ??
                        "anonymous";

                    return RateLimitPartition.GetFixedWindowLimiter(
                        partitionKey,
                        _ => CreateLimiterOptions(
                            rateLimitOptions.PermitLimit,
                            rateLimitOptions.QueueLimit,
                            rateLimitOptions.WindowSeconds));
                });
            }

            options.AddPolicy(
                AuthRateLimitPolicies.Password,
                httpContext => RateLimitPartition.GetFixedWindowLimiter(
                    AuthRateLimitPartitionKey.Resolve(httpContext, AuthRateLimitPolicies.Password),
                    _ => CreateLimiterOptions(
                        authRateLimitOptions.PasswordPermitLimit,
                        0,
                        authRateLimitOptions.PasswordWindowSeconds)));
            options.AddPolicy(
                AuthRateLimitPolicies.Sms,
                httpContext => RateLimitPartition.GetFixedWindowLimiter(
                    AuthRateLimitPartitionKey.Resolve(httpContext, AuthRateLimitPolicies.Sms),
                    _ => CreateLimiterOptions(
                        authRateLimitOptions.SmsPermitLimit,
                        0,
                        authRateLimitOptions.SmsWindowSeconds)));
            options.OnRejected = WriteRateLimitProblemAsync;
        });

        return services;
    }

    private static FixedWindowRateLimiterOptions CreateLimiterOptions(
        int permitLimit,
        int queueLimit,
        int windowSeconds) =>
        new()
        {
            AutoReplenishment = true,
            PermitLimit = permitLimit,
            QueueLimit = queueLimit,
            QueueProcessingOrder = QueueProcessingOrder.OldestFirst,
            Window = TimeSpan.FromSeconds(windowSeconds)
        };

    private static async ValueTask WriteRateLimitProblemAsync(
        OnRejectedContext context,
        CancellationToken cancellationToken)
    {
        if (context.Lease.TryGetMetadata(MetadataName.RetryAfter, out var retryAfter))
        {
            context.HttpContext.Response.Headers.RetryAfter = ((int)retryAfter.TotalSeconds).ToString();
        }

        var problem = new ProblemDetails
        {
            Title = "Too many requests.",
            Status = StatusCodes.Status429TooManyRequests,
            Instance = context.HttpContext.Request.Path
        };
        problem.Extensions["code"] = "rate_limited";
        problem.Extensions["traceId"] = context.HttpContext.TraceIdentifier;

        context.HttpContext.Response.StatusCode = StatusCodes.Status429TooManyRequests;
        await context.HttpContext.Response.WriteAsJsonAsync(problem, cancellationToken);
    }
}
