use core::{fmt, mem::MaybeUninit};

use alloc::collections::{btree_map::Entry, BTreeMap};
use kernel_util::mem::PageBox;
use yggdrasil_abi::error::Error;

// TODO LRU
pub struct BlockCache<K: Ord + Eq> {
    table: BTreeMap<K, CachedBlock>,
    block_size: usize,
}

pub type CachedBlock = PageBox<[u8]>;

trait TryGetOrInsertWith<K, V> {
    fn try_get_or_insert_with<E, F: FnOnce() -> Result<V, E>>(
        &mut self,
        key: K,
        f: F,
    ) -> Result<&mut V, E>;
}

impl<K: Ord + Eq> BlockCache<K> {
    pub fn new(block_size: usize) -> Self {
        log::debug!("Block cache created with bs={}", block_size);
        Self {
            table: BTreeMap::new(),
            block_size,
        }
    }

    pub fn get_or_fetch_with<
        E: From<Error>,
        F: FnOnce(&mut PageBox<[MaybeUninit<u8>]>) -> Result<(), E>,
    >(
        &mut self,
        index: K,
        f: F,
    ) -> Result<&mut CachedBlock, E>
    where
        K: Copy + fmt::Display,
    {
        if let Entry::Vacant(entry) = self.table.entry(index) {
            let mut block = PageBox::new_uninit_slice(self.block_size)?;
            log::debug!("Missed block with index {}, fetching", index);
            f(&mut block)?;
            entry.insert(unsafe { block.assume_init_slice() });
        }

        Ok(self.table.get_mut(&index).unwrap())
    }

    pub fn evict(&mut self, index: K) {
        self.table.remove(&index);
    }
}