use core::{
    borrow::Borrow,
    fmt::Debug,
    hint::unlikely,
    ops::{Deref, Index},
};

use bit_field::BitField;

use crate::{
    memory::{PhyAddr, VirtAddr, VirtAddrTranslationExt},
    x86_64::{VirtAddrExt, registers::Cr4},
};

#[repr(C, align(4096))]
pub struct PageTable {
    entries: [PageTableEntry; 512],
}

impl PageTable {
    pub fn get(&self, index: u16) -> Option<PageTableEntry> {
        if index < 512 {
            let entry = self.entries[index as usize];

            if entry.present() { Some(entry) } else { None }
        } else {
            None
        }
    }
    pub fn get_unchecked(&self, index: u16) -> PageTableEntry {
        assert!(index < 512, "Page table index out of bounds");
        self.entries[index as usize]
    }
}

impl Index<u16> for PageTable {
    type Output = PageTableEntry;

    fn index(&self, index: u16) -> &Self::Output {
        assert!(index < 512, "Page table index out of bounds");
        &self.entries[index as usize]
    }
}

#[repr(transparent)]
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct PageTableEntry(PageTableEntryFlags);

impl Debug for PageTableEntry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("PageTableEntry")
            .field(
                "flags",
                &PageTableEntryFlags::from_bits_truncate(self.as_raw()),
            )
            .field("phy", &self.phy())
            .field("pk", &self.pk())
            .field("pat_index", &self.pat_index())
            .field("free_bits", &format_args!("{:#b}", self.free_bits()))
            .finish()
    }
}

impl AsRef<PageTableEntryFlags> for PageTableEntry {
    fn as_ref(&self) -> &PageTableEntryFlags {
        &self.0
    }
}

impl Borrow<PageTableEntryFlags> for PageTableEntry {
    fn borrow(&self) -> &PageTableEntryFlags {
        &self.0
    }
}

impl Deref for PageTableEntry {
    type Target = PageTableEntryFlags;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

bitflags::bitflags! {
    #[repr(transparent)]
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct PageTableEntryFlags: u64 {
        const PRESENT = 1 << 0;
        const WRITABLE = 1 << 1;
        const USER_ACCESSIBLE = 1 << 2;
        const WRITE_THROUGH = 1 << 3; // PWT
        const NO_CACHE = 1 << 4; // PCD
        const ACCESSED = 1 << 5;
        const DIRTY = 1 << 6;
        const HUGE_PAGE = 1 << 7;
        const GLOBAL = 1 << 8;
        const NO_EXECUTE = 1 << 63;
        const PAT_4K = 1 << 7;
        const PAT_HUGE = 1 << 12;
    }
}

impl PageTableEntry {
    pub fn from_raw(bits: u64) -> Self {
        Self(PageTableEntryFlags::from_bits_retain(bits))
    }
    pub fn as_raw(&self) -> u64 {
        self.0.bits()
    }
    pub fn as_mut_raw(&mut self) -> &mut u64 {
        self.0.0.bits_mut()
    }
    pub fn present(&self) -> bool {
        self.contains(PageTableEntryFlags::PRESENT)
    }
    pub fn phy(&self) -> PhyAddr {
        PhyAddr(self.as_raw().get_bits(12..52) << 12)
    }
    pub fn try_as_page_table(&self) -> Option<&PageTable> {
        if !self.contains(PageTableEntryFlags::PRESENT) {
            return None;
        }
        if self.contains(PageTableEntryFlags::HUGE_PAGE) {
            return None;
        }

        unsafe {
            Some(
                self.phy()
                    .as_hhdm_virt()
                    .as_ptr::<PageTable>()
                    .as_ref()
                    .unwrap_unchecked(),
            )
        }
    }
    pub unsafe fn as_page_table(&self) -> &PageTable {
        // entry must be present
        assert!(self.contains(PageTableEntryFlags::PRESENT));
        // if the entry is a huge page, it does not point to a deeper page table.
        assert!(!self.contains(PageTableEntryFlags::HUGE_PAGE));

        unsafe {
            self.phy()
                .as_hhdm_virt()
                .as_ptr::<PageTable>()
                .as_ref()
                .unwrap_unchecked()
        }
    }
    pub fn pk(&self) -> u8 {
        self.as_raw().get_bits(59..63) as u8
    }
    pub fn pat_index(&self) -> u8 {
        let pat = if self.contains(PageTableEntryFlags::HUGE_PAGE) {
            self.as_raw().get_bit(12) as u8
        } else {
            self.as_raw().get_bit(7) as u8
        };

        self.as_raw().get_bits(3..=4) as u8 | (pat << 2)
    }

    /// Returns the free bits in the page table entry, which are bits 9-11 and
    /// 52-58, combined into a single 10-bit value.
    pub fn free_bits(&self) -> u16 {
        let low = self.bits().get_bits(9..12) as u16;
        let high = self.bits().get_bits(52..59) as u16;
        low | (high << 3)
    }
    pub fn set_free_bits(&mut self, value: u16) {
        let low = (value & 0b111) as u64;
        let high = (value >> 3) as u64;
        self.as_mut_raw().set_bits(9..12, low);
        self.as_mut_raw().set_bits(52..59, high);
    }
}

impl VirtAddrTranslationExt for VirtAddr {
    fn into_phy_addr(self) -> Option<PhyAddr> {
        get_physical_addr(self)
    }
}

pub fn get_physical_addr(virt: VirtAddr) -> Option<PhyAddr> {
    let cr3 = crate::x86_64::registers::Cr3::read();
    let cr4 = crate::x86_64::registers::Cr4::read();

    let frame = cr3.phy();
    let page_table = unsafe {
        frame
            .as_hhdm_virt()
            .as_ptr::<PageTable>()
            .as_ref()
            .unwrap_unchecked()
    };

    let pml4_entry = if unlikely(cr4.contains(Cr4::LA57)) {
        let l5_entry = page_table[virt.page_table_index::<{ VirtAddr::PML5 }>()];

        l5_entry
            .try_as_page_table()?
            .get_unchecked(virt.page_table_index::<{ VirtAddr::PML4 }>())
    } else {
        page_table[virt.page_table_index::<{ VirtAddr::PML4 }>()]
    };

    let pdpt_entry = pml4_entry.try_as_page_table()?[virt.page_table_index::<{ VirtAddr::PDPT }>()];

    if pdpt_entry.contains(PageTableEntryFlags::HUGE_PAGE) {
        const PHY_MASK_1G: u64 = !((1 << 30) - 1);
        let phys_addr = (pdpt_entry.phy().0 & PHY_MASK_1G) + virt.offset_1g() as u64;
        return Some(PhyAddr(phys_addr));
    }

    let pd_entry = pdpt_entry.try_as_page_table()?[virt.page_table_index::<{ VirtAddr::PD }>()];

    if pd_entry.contains(PageTableEntryFlags::HUGE_PAGE) {
        const PHY_MASK_2M: u64 = !((1 << 21) - 1);
        let phys_addr = (pd_entry.phy().0 & PHY_MASK_2M) + virt.offset_2m() as u64;
        return Some(PhyAddr(phys_addr));
    }

    let pt_entry = pd_entry.try_as_page_table()?[virt.page_table_index::<{ VirtAddr::PT }>()];

    if !pt_entry.contains(PageTableEntryFlags::PRESENT) {
        return None;
    }

    let phy_addr = pt_entry.phy().as_hhdm_virt().0 + virt.offset_4k() as u64;

    Some(PhyAddr(phy_addr))
}