use libk_util::sync::spin_rwlock::IrqSafeRwLock;

use crate::error::UsbError;

pub struct UsbAddressAllocator {
    // 256 bits
    bitmap: IrqSafeRwLock<Bitmap>,
}

struct Bitmap {
    data: [u64; 4],
}

impl Bitmap {
    const fn bit_index(bit: u8) -> usize {
        (bit / 64) as usize
    }

    const fn bit_mask(bit: u8) -> u64 {
        1 << (bit % 64)
    }

    pub const fn new() -> Self {
        // First is 1 to prevent address 0 from being allocated
        Self { data: [1, 0, 0, 0] }
    }

    pub fn is_set(&self, bit: u8) -> bool {
        self.data[Self::bit_index(bit)] & Self::bit_mask(bit) != 0
    }

    pub fn set(&mut self, bit: u8) {
        self.data[Self::bit_index(bit)] |= Self::bit_mask(bit);
    }

    pub fn clear(&mut self, bit: u8) {
        self.data[Self::bit_index(bit)] &= !Self::bit_mask(bit);
    }
}

impl UsbAddressAllocator {
    pub fn new() -> Self {
        Self {
            bitmap: IrqSafeRwLock::new(Bitmap::new()),
        }
    }

    pub fn allocate(&self) -> Result<u8, UsbError> {
        let mut bitmap = self.bitmap.write();
        for bit in 0..=255 {
            if !bitmap.is_set(bit) {
                bitmap.set(bit);
                return Ok(bit);
            }
        }
        Err(UsbError::OutOfAddresses)
    }

    pub fn free(&self, address: u8) {
        let mut bitmap = self.bitmap.write();
        assert!(bitmap.is_set(address));
        bitmap.clear(address);
    }
}
