use crate::backend::memory::{PhysAddr, VirtAddr};

const TLB_SIZE: usize = 256;
const TLB_MASK: u32 = (TLB_SIZE - 1) as u32;

pub struct TLB {
    entries: [TlbEntry; TLB_SIZE],
}

#[derive(Debug, Clone, Copy)]
struct TlbEntry {
    virtual_address: VirtAddr,
    physical_address: PhysAddr,
    valid: bool,
}

impl TLB {
    pub fn new() -> Self {
        TLB {
            entries: [TlbEntry {
                virtual_address: 0,
                physical_address: 0,
                valid: false,
            }; TLB_SIZE],
        }
    }

    #[inline(always)]
    fn index(page: u32) -> usize {
        let hash = page ^ (page >> 12);
        (hash & TLB_MASK) as usize
    }

    #[inline(always)]
    pub fn lookup(&self, virtual_address: VirtAddr) -> Option<PhysAddr> {
        let page = virtual_address >> 12; // for 4KB pages
        let entry = self.entries[Self::index(page)];
        if entry.valid && entry.virtual_address == virtual_address {
            Some(entry.physical_address)
        } else {
            None
        }
    }

    #[inline(always)]
    pub fn insert(&mut self, virt: VirtAddr, phys: PhysAddr) {
        let idx = Self::index(virt) as usize;
        self.entries[idx] = TlbEntry {
            virtual_address: virt,
            physical_address: phys,
            valid: true,
        };
    }
}
