use crate::arch::PAGE_SIZE;
use contracts::requires;
use core::{arch::asm, fmt, str};

/// The number of bits looked up in each page table entry.
pub const PAGE_TABLE_BITS: usize = 9;

/// The number of levels of page tables.
pub const PAGE_TABLE_LEVELS: usize = 3;

/// An address space ID.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct ASID(u16);

impl ASID {
    /// The kernel's ASID.
    pub const KERNEL: ASID = ASID(0);
}

/// A single page table.
#[derive(Debug)]
#[repr(align(4096))]
pub struct PageTable([PageTableEntry; 512]);

impl PageTable {
    /// Set this as the root page table. Note that this does _not_ perform a TLB shootdown.
    ///
    /// # Safety
    ///
    /// - All the safety conditions that would apply for setting `satp` and issuing an
    ///   `sfence.vma`.
    #[requires((asid.0 & !0xfff) == 0)]
    #[requires(((self as *const PageTable as usize) & 0xff00_0000_0000_0fff) == 0)]
    #[inline(never)]
    pub unsafe fn make_current(&self, asid: ASID) {
        let mode = 8; // Sv39
        let addr = self as *const PageTable as usize as u64;
        let satp = (mode << 60) | ((asid.0 as u64) << 44) | (addr >> 12);
        asm!("sfence.vma", "csrw satp, {satp}", "sfence.vma", satp = in(reg) satp)
    }

    /// Iterates over shared references to the entries in this page table.
    pub fn iter(&self) -> impl Iterator<Item = &PageTableEntry> {
        self.0.iter()
    }

    /// Iterates over exclusive references to the entries in this page table.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut PageTableEntry> {
        self.0.iter_mut()
    }
}

/// An entry in a page table.
#[derive(Clone, Copy, Default, Eq, PartialEq)]
pub struct PageTableEntry(u64);

impl PageTableEntry {
    /// Returns the physical page number of the backing page or next level page table.
    #[requires(self.valid())]
    #[ensures((ret & !0x0000_0fff_ffff_ffff) == 0)]
    fn ppn(&self) -> u64 {
        (self.0 >> 10) & 0x0000_0fff_ffff_ffff
    }

    /// Returns the physical address of the backing page or next level page table.
    #[requires(self.valid())]
    #[ensures((ret & !0x003f_ffff_ffff_fc00) == 0)]
    pub fn addr(&self) -> u64 {
        self.ppn() << PAGE_TABLE_BITS
    }

    /// Returns a pointer to the backing page.
    #[requires(self.valid())]
    #[requires(self.leaf_pte())]
    #[ensures((ret as usize & !0x003f_ffff_ffff_fc00) == 0)]
    pub fn page(&self) -> *mut [u8; PAGE_SIZE] {
        self.addr() as *mut [u8; PAGE_SIZE]
    }

    /// Returns a pointer to the backing page table.
    #[requires(self.valid())]
    #[requires(!self.leaf_pte())]
    #[ensures((ret as usize & !0x003f_ffff_ffff_fc00) == 0)]
    pub fn page_table(&self) -> *mut PageTable {
        self.addr() as *mut PageTable
    }

    /// Sets the physical address of the backing page or next level page table.
    #[requires(self.valid())]
    #[requires((addr & !0x003f_ffff_ffff_fc00) == 0)]
    pub fn set_addr(&mut self, addr: u64) -> &mut PageTableEntry {
        let ppn = addr >> 12;
        self.0 &= !0x003f_ffff_ffff_fc00;
        self.0 |= ppn << 10;
        self
    }

    /// Returns whether the dirty bit is set.
    #[requires(self.valid())]
    pub fn dirty(&self) -> bool {
        (self.0 & (1 << 7)) != 0
    }

    /// Sets the dirty bit.
    #[requires(self.valid())]
    pub fn set_dirty(&mut self, dirty: bool) -> &mut PageTableEntry {
        self.0 &= !0b10000000;
        self.0 |= (dirty as u64) << 7;
        self
    }

    /// Returns whether the accessed bit is set.
    #[requires(self.valid())]
    pub fn accessed(&self) -> bool {
        (self.0 & (1 << 6)) != 0
    }

    /// Sets the accessed bit.
    #[requires(self.valid())]
    pub fn set_accessed(&mut self, accessed: bool) -> &mut PageTableEntry {
        self.0 &= !0b01000000;
        self.0 |= (accessed as u64) << 6;
        self
    }

    /// Returns whether the global bit is set.
    #[requires(self.valid())]
    pub fn global(&self) -> bool {
        (self.0 & (1 << 5)) != 0
    }

    /// Sets the global bit.
    #[requires(self.valid())]
    pub fn set_global(&mut self, global: bool) -> &mut PageTableEntry {
        self.0 &= !0b00100000;
        self.0 |= (global as u64) << 5;
        self
    }

    /// Returns whether the user bit is set.
    #[requires(self.valid())]
    pub fn user(&self) -> bool {
        (self.0 & (1 << 4)) != 0
    }

    /// Sets the user bit.
    #[requires(self.valid())]
    pub fn set_user(&mut self, user: bool) -> &mut PageTableEntry {
        self.0 &= !0b00010000;
        self.0 |= (user as u64) << 4;
        self
    }

    /// Returns whether the executable bit is set.
    #[requires(self.valid())]
    pub fn executable(&self) -> bool {
        (self.0 & (1 << 3)) != 0
    }

    /// Sets the executable bit.
    #[requires(self.valid())]
    pub fn set_executable(&mut self, executable: bool) -> &mut PageTableEntry {
        self.0 &= !0b00001000;
        self.0 |= (executable as u64) << 3;
        self
    }

    /// Returns whether the writable bit is set.
    #[requires(self.valid())]
    pub fn writable(&self) -> bool {
        (self.0 & (1 << 2)) != 0
    }

    /// Sets the writable bit.
    #[requires(self.valid())]
    pub fn set_writable(&mut self, writable: bool) -> &mut PageTableEntry {
        self.0 &= !0b00000100;
        self.0 |= (writable as u64) << 2;
        self
    }

    /// Returns whether the readable bit is set.
    #[requires(self.valid())]
    pub fn readable(&self) -> bool {
        (self.0 & (1 << 1)) != 0
    }

    /// Sets the readable bit.
    #[requires(self.valid())]
    pub fn set_readable(&mut self, readable: bool) -> &mut PageTableEntry {
        self.0 &= !0b00000010;
        self.0 |= (readable as u64) << 1;
        self
    }

    /// Returns whether the page table entry is for a leaf PTE.
    #[requires(self.valid())]
    pub fn leaf_pte(&self) -> bool {
        (self.0 & 0b1110) != 0
    }

    /// Sets the readable, writable, and executable bits at once.
    #[requires(self.valid())]
    pub fn set_rwx(
        &mut self,
        readable: bool,
        writable: bool,
        executable: bool,
    ) -> &mut PageTableEntry {
        self.set_readable(readable)
            .set_writable(writable)
            .set_executable(executable)
    }

    /// Returns whether the valid bit is set.
    pub fn valid(&self) -> bool {
        (self.0 & (1 << 0)) != 0
    }

    /// Sets the valid bit.
    pub fn set_valid(&mut self, valid: bool) -> &mut PageTableEntry {
        self.0 &= !0b00000001;
        self.0 |= valid as u64;
        self
    }
}

impl fmt::Debug for PageTableEntry {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        if self.valid() {
            let mut flags = *b"DAGUXWRV";
            if !self.dirty() {
                flags[0] = b' ';
            }
            if !self.accessed() {
                flags[1] = b' ';
            }
            if !self.global() {
                flags[2] = b' ';
            }
            if !self.user() {
                flags[3] = b' ';
            }
            if !self.executable() {
                flags[4] = b' ';
            }
            if !self.writable() {
                flags[5] = b' ';
            }
            if !self.readable() {
                flags[6] = b' ';
            }

            // UNWRAP: The flags must be ASCII.
            let addr = self.addr() as *const ();
            let flags = str::from_utf8(&flags).unwrap();
            write!(fmt, "PageTableEntry({addr:018p}, {flags})")
        } else {
            write!(fmt, "PageTableEntry({:#018x}, INVALID)", self.0)
        }
    }
}