use core::ops::DerefMut;
use core::ops::Index;
use core::ptr::addr_of;

use intbits::Bits;
use x86_64::PhysAddr;
use x86_64::VirtAddr;
use x86_64::registers::control;
use x86_64::registers::control::Cr4Flags;
use x86_64::structures::paging::PageTable;
use x86_64::structures::paging::PageTableFlags;
use x86_64::structures::paging::page_table::PageTableEntry;

use crate::log::*;
use crate::memory::HHDM_RESPONSE;
use crate::memory::format;

#[allow(dead_code)]
pub fn map_page(virt: VirtAddr, mapping: PageTableEntry) -> Option<VirtAddr> {
    // TODO:
    // This function should create necessary page tables until the
    // supplied virtual address maps to the supplied physical address.
    // Currently, it *will* give up if there are any missing tables on the path.
    // It needs to be altered to support creating new intermediary tables.
    // This will likely not be a necessary function in the future if VMM is left to processes.
    log_info!(
        "Mapping 0x{:x} to 0x{:x} with flags {:?}",
        virt,
        mapping.addr(),
        mapping.flags()
    );
    let root_table = unsafe { active_root_table() };
    let mut current_table = root_table;
    let table_indices = [
        virt.as_u64().bits(0..=11) as usize,  // Page offset
        virt.as_u64().bits(12..=20) as usize, // Page Table index
        virt.as_u64().bits(21..=29) as usize, // Page Directory index
        virt.as_u64().bits(30..=38) as usize, // Page Directory Pointer Table index
        virt.as_u64().bits(39..=47) as usize, // Page Map Level 4 index
        virt.as_u64().bits(48..=56) as usize, // Page Map Level 5 index
    ];
    while current_table.1 > 0 {
        log_trace!(
            "Searching through a level {} table at 0x{:x}",
            current_table.1,
            addr_of!(*current_table.0) as usize
        );
        log_trace!("Index: {}", table_indices[current_table.1]);
        let current_table_entry: &mut PageTableEntry =
            &mut current_table.0.deref_mut()[table_indices[current_table.1]];
        if current_table_entry
            .flags()
            .contains(PageTableFlags::PRESENT)
            && !current_table_entry.is_unused()
        {
            if current_table_entry
                .flags()
                .contains(PageTableFlags::HUGE_PAGE)
            {
                if current_table.1 == 3 {
                    // Huge page (1GiB), last 30 bits preserved
                    log_error!("Could not map: Entry on path is Huge Page!");
                    return None;
                } else if current_table.1 == 2 {
                    // Large page (2MiB), last 21 bits preserved
                    log_error!("Could not map: Entry on path is Large Page!");
                    return None;
                }
            } else if current_table.1 > 1 {
                current_table = (
                    unsafe { &mut *phys_to_virt(current_table_entry.addr()).as_mut_ptr() }
                        as &mut PageTable,
                    current_table.1 - 1,
                );
            } else {
                current_table_entry.set_addr(mapping.addr(), mapping.flags());
                return Some(virt);
            }
        } else {
            log_error!(
                "Could not map: Level {} entry on path is not present!",
                current_table.1
            );
            return None;
        }
    }
    // Current table is level 1
    unreachable!()
}

#[allow(dead_code)]
pub fn unmap_page(_virt: VirtAddr) -> Result<(), ()> {
    // TODO:
    // This funciton should remove any present mappings
    // corresponding to the supplied virtual address.
    todo!()
}

#[allow(dead_code)]
/// Safety: CR3 must never be modified while a PageTable returned from
/// this function is in scope. Otherwise, it could point to stale or even
/// invalid tables.
pub unsafe fn active_root_table() -> (&'static mut PageTable, usize) {
    use x86_64::registers::control::Cr3;
    let root_table_frame = Cr3::read().0;
    let virt = phys_to_virt(root_table_frame.start_address());
    let page_table_ptr: *mut PageTable = virt.as_mut_ptr();

    (
        unsafe { &mut *page_table_ptr },
        match control::Cr4::read().contains(Cr4Flags::L5_PAGING) {
            true => 5,
            false => 4,
        },
    )
}

pub fn walk_tables() {
    iter_table(
        match control::Cr4::read().contains(Cr4Flags::L5_PAGING) {
            true => 5,
            false => 4,
        },
        unsafe { active_root_table().0 },
    );
}

#[allow(dead_code)]
pub fn iter_table(level: usize, table: &PageTable) {
    for (i, entry) in table.iter().enumerate() {
        if !entry.is_unused() & entry.flags().contains(PageTableFlags::PRESENT) {
            log_trace!("L{} Entry {}: {:?}", level, i, entry);
            if (level > 1) & !entry.flags().contains(PageTableFlags::HUGE_PAGE) {
                let phys = entry.addr();
                let virt = phys_to_virt(phys).as_mut_ptr();
                let new_table: &PageTable = unsafe { &*virt };
                iter_table(level - 1, new_table);
            }
        }
    }
}

pub fn phys_to_virt(phys: PhysAddr) -> VirtAddr {
    VirtAddr::new(phys.as_u64() + HHDM_RESPONSE.offset)
}

#[allow(dead_code)]
pub fn virt_to_phys(virt: VirtAddr) -> Option<PhysAddr> {
    let root_table = unsafe { active_root_table() };
    let mut current_table = root_table;
    let table_indices = [
        virt.as_u64().bits(0..=11) as usize,  // Page offset
        virt.as_u64().bits(12..=20) as usize, // Page Table index
        virt.as_u64().bits(21..=29) as usize, // Page Directory index
        virt.as_u64().bits(30..=38) as usize, // Page Directory Pointer Table index
        virt.as_u64().bits(39..=47) as usize, // Page Map Level 4 index
        virt.as_u64().bits(48..=56) as usize, // Page Map Level 5 index
    ];
    while current_table.1 > 0 {
        log_trace!(
            "Searching through a level {} table at 0x{:x}",
            current_table.1,
            addr_of!(*current_table.0) as usize
        );
        log_trace!("Index: {}", table_indices[current_table.1]);
        let current_table_entry = current_table.0.index(table_indices[current_table.1]);
        if current_table_entry
            .flags()
            .contains(PageTableFlags::PRESENT)
            && !current_table_entry.is_unused()
        {
            if current_table_entry
                .flags()
                .contains(PageTableFlags::HUGE_PAGE)
            {
                if current_table.1 == 3 {
                    // Huge page (1GiB), last 30 bits preserved
                    return Some(PhysAddr::new(
                        current_table_entry.addr().as_u64() | virt.as_u64().bits(0..=29),
                    ));
                } else if current_table.1 == 2 {
                    // Large page (2MiB), last 21 bits preserved
                    return Some(PhysAddr::new(
                        current_table_entry.addr().as_u64() | virt.as_u64().bits(0..=20),
                    ));
                }
            } else if current_table.1 > 1 {
                current_table = (
                    unsafe { &mut *phys_to_virt(current_table_entry.addr()).as_mut_ptr() }
                        as &mut PageTable,
                    current_table.1 - 1,
                );
            } else {
                return Some(PhysAddr::new(
                    current_table_entry.addr().as_u64() | virt.as_u64().bits(0..=11),
                ));
            }
        } else {
            return None;
        }
    }
    // Current table is level 1
    unreachable!()
}
