use crate::block_device::{BlockDevice, BLOCK_SIZE};
use std::sync::Arc;

pub struct Bitmap {
    pub starting_block: usize,
    pub length: usize,
    pub device: Arc<dyn BlockDevice>,
    pub data: Vec<u8>,
    pub count: u64,
}

impl Bitmap {
    pub(crate) fn new(starting_block: usize, length: usize, device: Arc<dyn BlockDevice>, count: u64) -> Self {
        Self {
            starting_block,
            length,
            device,
            data: vec![0u8; length * BLOCK_SIZE],
            count,
        }
    }

    pub(crate) fn load(starting_block: usize, length: usize, device: Arc<dyn BlockDevice>, count: u64) -> Self {
        let mut data = vec![0u8; length * BLOCK_SIZE];
        for id in 0 .. length {
            let block_id = starting_block + id;
            device.read(block_id, &mut data[id * BLOCK_SIZE .. (id + 1) * BLOCK_SIZE]);
        }
        Self {
            starting_block,
            length,
            device,
            data,
            count,
        }
    }

    pub(crate) fn write_back(&self) {
        for id in 0 .. self.length {
            let block_id = self.starting_block + id;
            self.device.write(block_id, &self.data[id * BLOCK_SIZE .. (id + 1) * BLOCK_SIZE]);
        }
    }

    // Allocate, but doesn't modify self.
    pub(crate) fn peek(&self) -> Option<usize> {
        for (i, byte) in self.data.iter().enumerate() {
            let leading_ones = byte.leading_ones();
            if leading_ones != 8 {
                return Some(i * 8 + leading_ones as usize);
            }
        }
        None
    }

    pub(crate) fn allocate(&mut self) -> Option<usize> {
        for (i, byte) in self.data.iter_mut().enumerate() {
            let leading_ones = byte.leading_ones();
            if leading_ones != 8 {
                *byte |= (1 << (7 - leading_ones)) as u8;
                self.count += 1;
                return Some(i * 8 + leading_ones as usize);
            }
        }
        None
    }

    pub(crate) fn query(&self, index: usize) -> bool {
        if index == 0 {
            false
        } else {
            self.data[index / 8] & ((1 << (7 - index % 8)) as u8) != 0
        }
    }

    pub(crate) fn deallocate(&mut self, index: usize) -> bool {
        if self.query(index) {
            let mask = !(1u8 << (7 - index % 8));
            self.count -= 1;
            self.data[index / 8] &= mask;
            true
        } else {
            false
        }
    }
}