use core::{
    marker::PhantomData,
    sync::atomic::{AtomicU16, Ordering},
};

use libk_mm_interface::{
    address::{AsPhysicalAddress, PhysicalAddress},
    pointer::PhysicalRefMut,
    process::{PageAttributeUpdate, ProcessAddressSpaceManager},
    table::{
        EntryLevel, EntryLevelDrop, EntryLevelExt, MapAttributes, NextPageTable, TableAllocator,
    },
};
use yggdrasil_abi::error::Error;

use crate::mem::{
    clone_kernel_tables,
    table::{PageAttributes, PageEntry},
};

use super::{
    table::{DroppableRange, PageTable, L1, L2, L3},
    KernelTableManagerImpl, USER_BOUNDARY,
};

pub struct ProcessAddressSpaceImpl<TA: TableAllocator> {
    l1: PhysicalRefMut<'static, PageTable<L1>, KernelTableManagerImpl>,
    asid: u16,
    _pd: PhantomData<TA>,
}

impl<TA: TableAllocator> ProcessAddressSpaceManager<TA> for ProcessAddressSpaceImpl<TA> {
    const LOWER_LIMIT_PFN: usize = 8;
    const UPPER_LIMIT_PFN: usize = (16 << 30) / L3::SIZE;

    fn new() -> Result<Self, Error> {
        static LAST_ASID: AtomicU16 = AtomicU16::new(1);

        let mut l1 = unsafe {
            PhysicalRefMut::<'static, PageTable<L1>, KernelTableManagerImpl>::map(
                TA::allocate_page_table()?,
            )
        };

        for i in 0..512 {
            l1[i] = PageEntry::INVALID;
        }
        // Copy the kernel mappings
        clone_kernel_tables(&mut l1);

        let asid = LAST_ASID.fetch_add(1, Ordering::AcqRel);

        Ok(Self {
            l1,
            asid,
            _pd: PhantomData,
        })
    }

    unsafe fn map_page(
        &mut self,
        address: usize,
        physical: PhysicalAddress,
        flags: MapAttributes,
    ) -> Result<(), Error> {
        self.write_l3_entry(
            address,
            PageEntry::page(physical, to_page_attributes(flags)),
            false,
        )
        .unwrap();
        Ok(())
    }

    unsafe fn update_page_attributes(
        &mut self,
        address: usize,
        update: &PageAttributeUpdate,
    ) -> Result<(), Error> {
        self.update_l3_entry(address, |entry| entry.update(update))
    }

    unsafe fn unmap_page(&mut self, address: usize) -> Result<(PhysicalAddress, bool), Error> {
        self.pop_l3_entry(address)
    }

    fn translate(&self, address: usize) -> Result<(PhysicalAddress, MapAttributes), Error> {
        self.read_l3_entry(address).ok_or(Error::DoesNotExist)
    }

    fn as_address_with_asid(&self) -> (u64, u64) {
        let physical = unsafe { self.l1.as_physical_address() }.into_u64();
        (physical, self.asid as u64)
    }

    unsafe fn clear(&mut self) {
        unsafe { self.l1.drop_range::<TA>(L1::DROPPABLE_RANGE) };
    }
}

impl<TA: TableAllocator> ProcessAddressSpaceImpl<TA> {
    // Write a single 4KiB entry
    fn write_l3_entry(
        &mut self,
        virt: usize,
        entry: PageEntry<L3>,
        overwrite: bool,
    ) -> Result<(), Error> {
        if virt >= USER_BOUNDARY {
            log::warn!("Tried to map a userspace page to a non-userspace virtual region");
            return Err(Error::InvalidArgument);
        }

        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        let mut l2 = self.l1.get_mut_or_alloc::<TA>(l1i)?;
        let mut l3 = l2.get_mut_or_alloc::<TA>(l2i)?;

        if l3[l3i].is_present() && !overwrite {
            todo!();
        }

        l3[l3i] = entry;
        super::tlb_flush_va_asid(virt, self.asid as usize);

        Ok(())
    }

    fn update_l3_entry<F: FnOnce(&mut PageEntry<L3>) -> Result<(), Error>>(
        &mut self,
        virt: usize,
        mapper: F,
    ) -> Result<(), Error> {
        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        // TODO somehow drop tables if they're known to be empty?
        let mut l2 = self.l1.get_mut(l1i).ok_or(Error::DoesNotExist)?;
        let mut l3 = l2.get_mut(l2i).ok_or(Error::DoesNotExist)?;

        let entry = &mut l3[l3i];
        if !entry.is_present() {
            return Err(Error::DoesNotExist);
        }
        mapper(entry)?;
        super::tlb_flush_va_asid(virt, self.asid as usize);

        Ok(())
    }

    fn pop_l3_entry(&mut self, virt: usize) -> Result<(PhysicalAddress, bool), Error> {
        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        // TODO somehow drop tables if they're known to be empty?
        let mut l2 = self.l1.get_mut(l1i).ok_or(Error::DoesNotExist)?;
        let mut l3 = l2.get_mut(l2i).ok_or(Error::DoesNotExist)?;

        let entry = l3[l3i];
        let page = entry.as_page().ok_or(Error::DoesNotExist)?;
        let dirty = entry.is_dirty();

        l3[l3i] = PageEntry::INVALID;
        super::tlb_flush_va_asid(virt, self.asid as usize);

        Ok((page, dirty))
    }

    fn read_l3_entry(&self, virt: usize) -> Option<(PhysicalAddress, MapAttributes)> {
        if virt >= USER_BOUNDARY {
            log::warn!("Tried read an userspace page to a non-userspace virtual region");
            return None;
        }

        let l1i = virt.page_index::<L1>();
        let l2i = virt.page_index::<L2>();
        let l3i = virt.page_index::<L3>();

        let l2 = self.l1.get(l1i)?;
        let l3 = l2.get(l2i)?;

        let page = l3[l3i].as_page()?;

        Some((
            page.add(virt & 0xFFF),
            to_map_attributes(l3[l3i].attributes()),
        ))
    }
}

impl<TA: TableAllocator> Drop for ProcessAddressSpaceImpl<TA> {
    fn drop(&mut self) {
        // SAFETY: with safe usage of the ProcessAddressSpaceImpl, clearing and dropping
        // is safe, no one refers to the memory
        unsafe {
            self.clear();
            let l1_phys = self.l1.as_physical_address();
            TA::free_page_table(l1_phys);
            super::tlb_flush_asid(self.asid as usize);
        }
    }
}

fn to_page_attributes(src: MapAttributes) -> PageAttributes {
    let mut result = PageAttributes::R | PageAttributes::X;
    if src.contains(MapAttributes::USER_WRITE) {
        result |= PageAttributes::W;
    }
    if src.intersects(MapAttributes::USER_READ | MapAttributes::USER_WRITE) {
        result |= PageAttributes::U;
    }
    if src.contains(MapAttributes::DIRTY) {
        result |= PageAttributes::SW_DIRTY;
    }
    result
}

fn to_map_attributes(src: PageAttributes) -> MapAttributes {
    let mut result = MapAttributes::NON_GLOBAL;

    if src.contains(PageAttributes::U) {
        result |= MapAttributes::USER_READ;

        if src.contains(PageAttributes::W) {
            result |= MapAttributes::USER_WRITE;
        }
    }

    if src.contains(PageAttributes::SW_DIRTY) {
        result |= MapAttributes::DIRTY;
    }

    result
}
