use core::arch::asm;

use crate::memory::{PAGE_SIZE, PhysicalMemoryManager};

const PRESENT: u64 = 1 << 0;
const WRITABLE: u64 = 1 << 1;
const HUGE_PAGE: u64 = 1 << 7;
const CR4_PSE: u64 = 1 << 4;

const ENTRY_COUNT: usize = 512;
const IDENTITY_MAP_GIB: usize = 512;
const PDPTE_COUNT: usize = IDENTITY_MAP_GIB / 2;

#[derive(Clone, Copy)]
pub struct PageFlags(u64);

impl PageFlags {
    #[allow(dead_code)]
    pub const PRESENT: Self = Self(PRESENT);
    pub const WRITABLE: Self = Self(WRITABLE);
    #[allow(dead_code)]
    pub const USER: Self = Self(1 << 2);
    #[allow(dead_code)]
    pub const NO_EXECUTE: Self = Self(1 << 63);

    pub const fn bits(self) -> u64 {
        self.0
    }

    #[allow(dead_code)]
    pub const fn union(self, other: Self) -> Self {
        Self(self.0 | other.0)
    }
}

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

impl PageTable {
    fn zero(&mut self) {
        self.entries.fill(0);
    }
}

pub struct AddressSpace {
    pml4: *mut PageTable,
    identity_pdpt: *mut PageTable,
    identity_pds: [*mut PageTable; PDPTE_COUNT],
}

impl AddressSpace {
    pub unsafe fn new_kernel(memory: &mut PhysicalMemoryManager) -> Option<Self> {
        let pml4 = memory.alloc_page()? as *mut PageTable;
        let identity_pdpt = memory.alloc_page()? as *mut PageTable;
        let mut identity_pds = [core::ptr::null_mut(); PDPTE_COUNT];

        let identity_flags = PRESENT | WRITABLE | PageFlags::USER.bits();

        unsafe {
            (*pml4).zero();
            (*identity_pdpt).zero();
            (*pml4).entries[0] = identity_pdpt as u64 | identity_flags;
        }

        for (pdpt_index, pd) in identity_pds.iter_mut().enumerate() {
            *pd = memory.alloc_page()? as *mut PageTable;

            unsafe {
                (**pd).zero();
                (*identity_pdpt).entries[pdpt_index] = *pd as u64 | identity_flags;

                for pd_index in 0..ENTRY_COUNT {
                    let physical = ((pdpt_index * ENTRY_COUNT + pd_index) as u64) * 2 * 1024 * 1024;
                    (**pd).entries[pd_index] = physical | identity_flags | HUGE_PAGE;
                }
            }
        }

        Some(Self {
            pml4,
            identity_pdpt,
            identity_pds,
        })
    }

    pub unsafe fn map_page(
        &mut self,
        memory: &mut PhysicalMemoryManager,
        virtual_address: u64,
        physical_address: u64,
        flags: PageFlags,
    ) -> Option<()> {
        if !is_aligned(virtual_address) || !is_aligned(physical_address) {
            return None;
        }

        let table = unsafe { self.walk_create(memory, virtual_address)? };
        let index = pt_index(virtual_address);

        unsafe {
            if (*table).entries[index] & PRESENT != 0 {
                return None;
            }
            (*table).entries[index] = physical_address | flags.bits() | PRESENT;
            flush_tlb_one(virtual_address);
        }

        Some(())
    }

    pub unsafe fn unmap_page(&mut self, virtual_address: u64) -> Option<u64> {
        let table = unsafe { self.walk(virtual_address)? };
        let index = pt_index(virtual_address);

        unsafe {
            let entry = (*table).entries[index];
            if entry & PRESENT == 0 {
                return None;
            }
            (*table).entries[index] = 0;
            flush_tlb_one(virtual_address);
            Some(entry & 0x000f_ffff_ffff_f000)
        }
    }

    pub unsafe fn translate(&self, virtual_address: u64) -> Option<u64> {
        let pml4 = unsafe { &*self.pml4 };
        let pml4e = pml4.entries[pml4_index(virtual_address)];
        let pdpt = next_table(pml4e)?;

        let pdpte = unsafe { (*pdpt).entries[pdpt_index(virtual_address)] };
        if pdpte & HUGE_PAGE != 0 {
            return Some((pdpte & 0x000f_ffff_c000_0000) | (virtual_address & 0x3fff_ffff));
        }
        let pd = next_table(pdpte)?;

        let pde = unsafe { (*pd).entries[pd_index(virtual_address)] };
        if pde & HUGE_PAGE != 0 {
            return Some((pde & 0x000f_ffff_ffe0_0000) | (virtual_address & 0x1f_ffff));
        }
        let pt = next_table(pde)?;

        let pte = unsafe { (*pt).entries[pt_index(virtual_address)] };
        if pte & PRESENT == 0 {
            None
        } else {
            Some((pte & 0x000f_ffff_ffff_f000) | (virtual_address & 0xfff))
        }
    }

    pub unsafe fn activate(&self) {
        unsafe {
            let cr4: u64;
            asm!("mov {}, cr4", out(reg) cr4, options(nostack, preserves_flags));
            asm!(
                "mov cr4, {}",
                in(reg) cr4 | CR4_PSE,
                options(nostack, preserves_flags)
            );
            asm!(
                "mov cr3, {}",
                in(reg) self.pml4 as u64,
                options(nostack, preserves_flags)
            );
        }
    }

    pub fn root_table(&self) -> u64 {
        self.pml4 as u64
    }

    pub fn identity_tables(&self) -> (*mut PageTable, [*mut PageTable; PDPTE_COUNT]) {
        (self.identity_pdpt, self.identity_pds)
    }

    unsafe fn walk_create(
        &mut self,
        memory: &mut PhysicalMemoryManager,
        virtual_address: u64,
    ) -> Option<*mut PageTable> {
        let mut table = self.pml4;
        for index in [
            pml4_index(virtual_address),
            pdpt_index(virtual_address),
            pd_index(virtual_address),
        ] {
            unsafe {
                let entry = &mut (*table).entries[index];
                if *entry & HUGE_PAGE != 0 {
                    return None;
                }
                if *entry & PRESENT == 0 {
                    let new_table = memory.alloc_page()? as *mut PageTable;
                    (*new_table).zero();
                    *entry = new_table as u64 | PRESENT | WRITABLE | PageFlags::USER.bits();
                }
                table = (*entry & 0x000f_ffff_ffff_f000) as *mut PageTable;
            }
        }

        Some(table)
    }

    unsafe fn walk(&self, virtual_address: u64) -> Option<*mut PageTable> {
        let mut table = self.pml4;
        for index in [
            pml4_index(virtual_address),
            pdpt_index(virtual_address),
            pd_index(virtual_address),
        ] {
            unsafe {
                let entry = (*table).entries[index];
                if entry & PRESENT == 0 || entry & HUGE_PAGE != 0 {
                    return None;
                }
                table = (entry & 0x000f_ffff_ffff_f000) as *mut PageTable;
            }
        }

        Some(table)
    }
}

fn next_table(entry: u64) -> Option<*mut PageTable> {
    if entry & PRESENT == 0 {
        None
    } else {
        Some((entry & 0x000f_ffff_ffff_f000) as *mut PageTable)
    }
}

fn pml4_index(address: u64) -> usize {
    ((address >> 39) & 0x1ff) as usize
}

fn pdpt_index(address: u64) -> usize {
    ((address >> 30) & 0x1ff) as usize
}

fn pd_index(address: u64) -> usize {
    ((address >> 21) & 0x1ff) as usize
}

fn pt_index(address: u64) -> usize {
    ((address >> 12) & 0x1ff) as usize
}

fn is_aligned(address: u64) -> bool {
    address as usize & (PAGE_SIZE - 1) == 0
}

unsafe fn flush_tlb_one(virtual_address: u64) {
    unsafe {
        asm!("invlpg [{}]", in(reg) virtual_address, options(nostack, preserves_flags));
    }
}