use crate::arch::{PAGE_SIZE, PAGE_SIZE_BITS};
use contracts::requires;
use core::{arch::asm, fmt, iter, ops::RangeInclusive, str};
use either::Either;
use vernos_utils::debug;

/// 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.
#[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()
    }
}

impl fmt::Debug for PageTable {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        // Get an iterator over the valid leaf page table entries.
        let mut mappings = iter_level_2_mappings(&self.0, 0).peekable();

        // Make an iterator that merges adjacent entries that have the same flags.
        let mappings = iter::from_fn(|| {
            let (entry, mut vaddrs) = mappings.next()?;
            let paddrs_start = entry.addr() as usize;
            let mut len = (vaddrs.end() - vaddrs.start()) + 1;

            while let Some((next_entry, next_vaddrs)) = mappings.peek() {
                let next_paddrs_start = next_entry.addr() as usize;

                if entry.flag_bits() != next_entry.flag_bits()
                    || vaddrs.end().wrapping_add(1) != *next_vaddrs.start()
                    || paddrs_start.wrapping_add(len) != next_paddrs_start
                {
                    break;
                }
                // UNWRAP: .peek() already showed us that there's a next entry.
                let (_, next_vaddrs) = mappings.next().unwrap();
                vaddrs = *vaddrs.start()..=*next_vaddrs.end();
                len = (next_vaddrs.end() - vaddrs.start()) + 1;
            }
            let paddrs = paddrs_start..=paddrs_start + (len - 1);
            Some((entry, vaddrs, paddrs))
        });

        // Turn the iterator into an iterator over Debugs.
        let debug_mappings = mappings.map(|(entry, vaddrs, paddrs)| {
            debug(move |fmt| {
                let flags = entry.flags_str();
                // UNWRAP: The flags must be ASCII by the postcondition of flags_str().
                let flags = str::from_utf8(&flags).unwrap();
                write!(
                    fmt,
                    "[V|{:16x}-{:16x}][P|{:16x}-{:16x}][F|{}]",
                    *vaddrs.start(),
                    *vaddrs.end(),
                    *paddrs.start(),
                    *paddrs.end(),
                    flags
                )
            })
        });

        fmt.debug_list().entries(debug_mappings).finish()
    }
}

/// 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 bits of the entry that correspond to flags.
    ///
    /// This isn't `pub` because this isn't portable, though maybe it makes sense to instead export
    /// a predicate for "do these two entries have the _same_ flags bits," since that should be
    /// more portable.
    #[requires(self.valid())]
    #[ensures((ret & !0xffc0_0000_0000_03ff) == 0)]
    fn flag_bits(&self) -> u64 {
        self.0 & 0xffc0_0000_0000_03ff
    }

    /// Returns bytes that correspond to an ASCII string with the flags.
    #[requires(self.valid())]
    #[ensures(ret.iter().all(|ch| ch.is_ascii()))]
    fn flags_str(&self) -> [u8; 7] {
        let mut flags = *b"rwxugad";
        let char_disabled = b'-';
        if !self.readable() {
            flags[0] = char_disabled;
        }
        if !self.writable() {
            flags[1] = char_disabled;
        }
        if !self.executable() {
            flags[2] = char_disabled;
        }
        if !self.user() {
            flags[3] = char_disabled;
        }
        if !self.global() {
            flags[4] = char_disabled;
        }
        if !self.accessed() {
            flags[5] = char_disabled;
        }
        if !self.dirty() {
            flags[6] = char_disabled;
        }
        flags
    }

    /// 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_SIZE_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 addr = self.addr() as *const ();
            let flags = self.flags_str();
            // UNWRAP: The flags must be ASCII by the postcondition of flags_str().
            let flags = str::from_utf8(&flags).unwrap();
            write!(fmt, "PageTableEntry({addr:018p}, {flags})")
        } else {
            write!(fmt, "PageTableEntry({:#018x}, INVALID)", self.0)
        }
    }
}

/// See `PageTable::iter_mappings`. This needs to be its own function because of `impl Trait`; we
/// can't allocate here, and we want a fixed-size iterator.
fn iter_level_2_mappings(
    table: &[PageTableEntry; 512],
    base_addr: usize,
) -> impl '_ + Iterator<Item = (PageTableEntry, RangeInclusive<usize>)> {
    const ENTRY_SIZE: usize = 1 << (12 + 9 + 9);
    table
        .iter()
        .enumerate()
        .filter(|(_, entry)| entry.valid())
        .flat_map(move |(i, &entry)| {
            let mut start_addr = base_addr + i * ENTRY_SIZE;
            if i >= 256 {
                start_addr += 0xffff_ff80_0000_0000;
            }
            if entry.leaf_pte() {
                Either::Left(iter::once((
                    entry,
                    start_addr..=start_addr + (ENTRY_SIZE - 1),
                )))
            } else {
                let next_table = unsafe { &(*entry.page_table()).0 };
                Either::Right(iter_level_1_mappings(next_table, start_addr))
            }
        })
}

/// See `PageTable::iter_mappings`. This needs to be its own function because of `impl Trait`; we
/// can't allocate here, and we want a fixed-size iterator.
fn iter_level_1_mappings(
    table: &[PageTableEntry; 512],
    base_addr: usize,
) -> impl '_ + Iterator<Item = (PageTableEntry, RangeInclusive<usize>)> {
    const ENTRY_SIZE: usize = 1 << (12 + 9);
    table
        .iter()
        .enumerate()
        .filter(|(_, entry)| entry.valid())
        .flat_map(move |(i, &entry)| {
            let start_addr = base_addr + i * ENTRY_SIZE;
            if entry.leaf_pte() {
                Either::Left(iter::once((
                    entry,
                    start_addr..=start_addr + (ENTRY_SIZE - 1),
                )))
            } else {
                let next_table = unsafe { &(*entry.page_table()).0 };
                Either::Right(iter_level_0_mappings(next_table, start_addr))
            }
        })
}

/// See `PageTable::iter_mappings`. This needs to be its own function because of `impl Trait`; we
/// can't allocate here, and we want a fixed-size iterator.
fn iter_level_0_mappings(
    table: &[PageTableEntry; 512],
    base_addr: usize,
) -> impl '_ + Iterator<Item = (PageTableEntry, RangeInclusive<usize>)> {
    const ENTRY_SIZE: usize = 1 << 12;
    table
        .iter()
        .enumerate()
        .filter(|(_, entry)| entry.valid())
        .map(move |(i, &entry)| {
            let start_addr = base_addr + i * ENTRY_SIZE;
            (entry, start_addr..=start_addr + (ENTRY_SIZE - 1))
        })
}