use crate::rough_panic;

#[cfg(feature = "arch_virt")]

unsafe extern "C" {
    fn _heap_start();
    fn _heap_size();
}

pub const BLOCK_SIZE: usize = 512;
#[cfg(feature = "arch_virt")]
pub const TOTAL_MEMORY: usize = 1048510;
#[cfg(feature = "arch_ppc32")]
pub const TOTAL_MEMORY: usize = 1024 * 1024; // 1MiB;
pub const MEM_BLOCKS: usize = TOTAL_MEMORY / BLOCK_SIZE;

pub struct MemoryManager {
    pub heap_start: usize,
    pub heap_size: usize,
    pub blockmap: [u8; (MEM_BLOCKS+7) / 8],
}

impl MemoryManager {
    #[cfg(feature = "arch_virt")]
    pub fn init() -> Self {
        Self {
            heap_start: _heap_start as _,
            heap_size: _heap_size as _,
            blockmap: [0; (MEM_BLOCKS+7) / 8],
        }
    }
    #[cfg(feature = "arch_ppc32")]
    pub fn init(heap_start: usize, heap_size: usize) -> Self {
        Self {
            heap_start: heap_start as _,
            heap_size: heap_size as _,
            blockmap: [0; (MEM_BLOCKS+7) / 8],
        }
    }

    pub fn block_to_addr(&self, block: usize) -> usize {
        block * BLOCK_SIZE + self.heap_start
    }
    
    pub fn addr_to_block(&self, addr: usize) -> usize {
        (addr - self.heap_start) / BLOCK_SIZE
    }
    
    pub fn alloc_one_block(&mut self) -> usize {
        /*
        for (i, v) in self.blockmap.iter_mut().enumerate() {
            for j in 0..8 {
                let block = i * 8 + j;
                let val = (1 << j) & *v;
                if val == 0 {
                    // free
                    *v |= 1 << j;
                    return block;
                }
            }
        }
        
        rough_panic(['o', 'o', 'm'])
         */
        self.alloc_n_blocks(1)
    }
    pub fn free_one_block(&mut self, block: usize) {
        self.blockmap[block / 8] &= !(1 << (block % 8));
    }
    
    // can easily fail if too many blocks are requested, will return 0 on failure
    pub fn alloc_n_blocks(&mut self, n: usize) -> usize {
        if n == 0 {
            return 0;
        }
        let mut first_block = None;
        let mut found = 0;
        for i in 0..self.blockmap.len() {
            for j in 0..8 {
                let block = i * 8 + j;
                let val = (1 << j) & self.blockmap[i];
                if val == 0 {
                    // this is free
                    self.blockmap[i] |= 1 << j;
                    if first_block.is_none() {
                        first_block = Some(block);
                    }
                    found += 1;
                    if found >= n {
                        return first_block.unwrap();
                    }
                } else {
                    // used, restart search
                    let mut i = 0;
                    while found > 0 {
                        found -= 1;
                        self.free_one_block(first_block.unwrap() + i);
                        i += 1;
                    }
                    first_block = None;
                }
            }
        }
        
        0
    }
    
    pub fn free_n_blocks(&mut self, block: usize, n: usize) {
        if n == 0 || block >= MEM_BLOCKS || block + n > MEM_BLOCKS {
            return;
        }
        for i in 0..n {
            self.free_one_block(block + i);
        }
    }
    
    pub fn used_blocks(&self) -> usize {
        let mut used_blocks = 0;
        for v in self.blockmap.iter() {
            for j in 0..8 {
                let val = (1 << j) & *v;
                if val != 0 {
                    // used
                    used_blocks += 1;
                }
            }
        }
        used_blocks
    }
}