using Polly;

namespace Foxcord.Rest.Rate;

public class DiscordResilienceStrategy(RateLimiter rateLimiter)
    : ResilienceStrategy<HttpResponseMessage>
{
    protected override async ValueTask<Outcome<HttpResponseMessage>> ExecuteCore<TState>(
        Func<ResilienceContext, TState, ValueTask<Outcome<HttpResponseMessage>>> callback, ResilienceContext context,
        TState state)
    {
        var path = context.Properties.GetValue(new ResiliencePropertyKey<string>("Path"), string.Empty);
        if (path == string.Empty) throw new DiscordRequestError("Path was not set in Polly context");

        var b = await rateLimiter.LockBucket(BucketKeyUtils.Parse(path), context.CancellationToken);

        var response = await callback(context, state).ConfigureAwait(context.ContinueOnCapturedContext);
        if (response.Exception is RateLimitError rateLimitError)
            b.Release(rateLimitError.Headers);
        else if (response.Result != null)
            b.Release(response.Result.Headers);

        return response;
    }
}

public class DiscordResilienceStrategyOptions : ResilienceStrategyOptions;

public static class DiscordResilienceStrategyExtensions
{
    public static ResiliencePipelineBuilder<HttpResponseMessage> AddDiscordStrategy(
        this ResiliencePipelineBuilder<HttpResponseMessage> builder, RateLimiter rateLimiter) =>
        builder.AddStrategy(_ => new DiscordResilienceStrategy(rateLimiter), new DiscordResilienceStrategyOptions());
}