using Microsoft.AspNetCore.RateLimiting;

using PrivaPub.Federation.Objects;
using PrivaPub.Federation.Signing;

using System.Threading.RateLimiting;

namespace PrivaPub.Infrastructure
{
	public static class RateLimiting
	{
		public const string Accounts = "accounts";
		public const string Inbox = "inbox";

		public static IServiceCollection PrivaPubRateLimiting(this IServiceCollection service) =>
			service.AddRateLimiter(options =>
			{
				options.RejectionStatusCode = StatusCodes.Status429TooManyRequests;
				options.AddPolicy(Accounts, context => RateLimitPartition.GetFixedWindowLimiter(
					context.Connection.RemoteIpAddress?.ToString() ?? "unknown",
					_ => new FixedWindowRateLimiterOptions { PermitLimit = 10, Window = TimeSpan.FromMinutes(1), QueueLimit = 0 }));
				options.AddPolicy(Inbox, context => RateLimitPartition.GetTokenBucketLimiter(
					SenderOrigin(context.Request) ?? "unsigned:" + context.Connection.RemoteIpAddress,
					_ => new TokenBucketRateLimiterOptions
					{
						TokenLimit = 300,
						TokensPerPeriod = 50,
						ReplenishmentPeriod = TimeSpan.FromSeconds(10),
						QueueLimit = 0
					}));
			});

		static string SenderOrigin(HttpRequest request)
		{
			var signature = request.Headers["Signature"].ToString();
			return string.IsNullOrEmpty(signature) ? default : Origin.Of(HttpSignatures.Parse(signature)?.KeyId);
		}
	}
}