// Copyright Marcus Del Favero 2025
// Licensed under the GNU AGPLv3 with an exception, see `README.md` for details
pub mod bucket;

use std::{collections::BTreeMap, net::SocketAddr, time::Instant};

use bucket::Bucket;

use super::ip_block::IpBlock;

use crate::net::serp::request::Request;

// Leaky bucket rate limiting: https://en.wikipedia.org/wiki/Leaky_bucket
pub struct RateLimiter {
	buckets: BTreeMap<(IpBlock, u32), Bucket>,
	next_gc: usize,
}

impl RateLimiter {
	pub fn new() -> RateLimiter {
		RateLimiter {
			buckets: BTreeMap::new(),
			next_gc: 0,
		}
	}

	/**
	 * Takes an address and the SERP request and returns if it's allowed to be
	 * performed or it has exceeded the rate limit.
	 */
	pub fn allowed(&mut self, addr: SocketAddr, request: &Request) -> bool {
		/*
		 * id: A unique value, doesn't matter what this is.
		 * to_add: Amount added to the bucket on allowed request.
		 * burst: Maximum capacity of bucket.
		 */
		let (id, to_add, burst) = match request {
			Request::Play(_) => (0, 4.0, 64.0),
			Request::Info => return true, // No rate limiting for Info requests as they're very cheap and pretty important to be available (AGPL requirement)
			Request::Discover => (1, 0.5, 24.0),
			Request::Publish(_) => (2, 16.0, 96.0),
		};

		let now = Instant::now();
		let bucket = self.buckets.entry((addr.into(), id)).or_insert_with(|| Bucket::new(now));

		let allowed = bucket.add(now, to_add, burst);
		if allowed {
			if self.next_gc == 0 {
				self.collect_garbage(now);
			} else {
				self.next_gc -= 1;
			}
		}
		allowed
	}

	fn collect_garbage(&mut self, now: Instant) {
		let prev_count = self.buckets.len();
		self.buckets.retain(|_, bucket| bucket.non_empty(now));
		let count = self.buckets.len();

		log::debug!("collected garbage, bucket count = {count} (previous count = {prev_count})");
		self.next_gc = count;
	}
}
