use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use tokio::sync::{OwnedSemaphorePermit, Semaphore};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HandlerLimitError;

#[derive(Clone)]
pub struct HandlerLimiter {
    sem: Arc<Semaphore>,
    max_inflight: usize,
    rejected: Arc<AtomicU64>,
}

impl HandlerLimiter {
    pub fn new(max_inflight: usize) -> Self {
        Self {
            sem: Arc::new(Semaphore::new(max_inflight)),
            max_inflight,
            rejected: Arc::new(AtomicU64::new(0)),
        }
    }

    pub fn try_acquire(&self) -> Result<OwnedSemaphorePermit, HandlerLimitError> {
        match self.sem.clone().try_acquire_owned() {
            Ok(permit) => Ok(permit),
            Err(_) => {
                self.rejected.fetch_add(1, Ordering::Relaxed);
                Err(HandlerLimitError)
            }
        }
    }

    pub fn inflight(&self) -> usize {
        self.max_inflight - self.sem.available_permits()
    }

    pub fn available(&self) -> usize {
        self.sem.available_permits()
    }

    pub fn rejected_total(&self) -> u64 {
        self.rejected.load(Ordering::Relaxed)
    }
}