#![cfg(target_arch = "x86_64")]

pub mod backtrace;
pub mod cpuid;
pub mod gdt;
pub mod idt;
pub mod instructions;
pub mod paging;
pub mod registers;

use core::{arch::asm, fmt::Debug};

pub use instructions::hlt;

use crate::memory::VirtAddr;

pub const PAGE_SIZE: usize = 4096;

pub fn halt_loop() -> ! {
    loop {
        hlt();
    }
}

#[repr(transparent)]
pub struct PortU16(pub u16);

impl PortU16 {
    pub fn offset(&self, offset: u16) -> Self {
        Self(self.0 + offset)
    }

    #[inline]
    pub unsafe fn read(&self) -> u16 {
        unsafe { instructions::read_u16(self.0) }
    }

    #[inline]
    pub unsafe fn write(&self, value: u16) {
        unsafe { instructions::write_u16(self.0, value) }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(C)]
// registers in the order they are pushed onto the stack:
pub struct Registers {
    pub r15: u64,
    pub r14: u64,
    pub r13: u64,
    pub r12: u64,
    pub r11: u64,
    pub r10: u64,
    pub r9: u64,
    pub r8: u64,
    pub rdi: u64,
    pub rsi: u64,
    pub rbp: u64,
    // rsp is not included here
    pub rbx: u64,
    pub rdx: u64,
    pub rcx: u64,
    pub rax: u64,
}

bitflags::bitflags! {
    #[repr(transparent)]
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct RFlags: u64 {
        const CARRY = 1 << 0;
        const DEFAULT = 1 << 1;
        const PARITY = 1 << 2;
        const AUXILIARY_CARRY = 1 << 4;
        const ZERO = 1 << 6;
        const SIGN = 1 << 7;
        const TRAP = 1 << 8;
        const INTERRUPT_ENABLE = 1 << 9;
        const DIRECTION = 1 << 10;
        const OVERFLOW = 1 << 11;
        const IOPL_LOW = 1 << 12;
        const IOPL_HIGH = 1 << 13;
        const NESTED_TASK = 1 << 14;
        const RESUME = 1 << 16;
        const VIRTUAL_8086_MODE = 1 << 17;
        const ALIGNMENT_CHECK = 1 << 18;
        const VIRTUAL_INTERRUPT_FLAG = 1 << 19;
        const VIRTUAL_INTERRUPT_PENDING = 1 << 20;
        const ID_FLAG = 1 << 21;
    }
}

pub trait VirtAddrExt {
    const LEVEL5: u8 = 4;
    const LEVEL4: u8 = 3;
    const LEVEL3: u8 = 2;
    const LEVEL2: u8 = 1;
    const LEVEL1: u8 = 0;

    const PT: u8 = Self::LEVEL1;
    const PD: u8 = Self::LEVEL2;
    const PDPT: u8 = Self::LEVEL3;
    const PML4: u8 = Self::LEVEL4;
    const PML5: u8 = Self::LEVEL5;

    fn is_canonical(&self) -> bool;
    fn page_table_index<const LEVEL: u8>(&self) -> u16;
    fn offset_4k(&self) -> u16;
    fn offset_2m(&self) -> u32;
    fn offset_1g(&self) -> u32;
    fn into_parts<E: sealed::PageSize>(self) -> (E::Entries, E::Offset);
}

pub struct L4Entries4K;
pub struct L4Entries2M;
pub struct L4Entries1G;

pub struct U12(pub u16);
pub struct U21(pub u32);
pub struct U30(pub u32);

pub(crate) mod sealed {
    use super::{L4Entries1G, L4Entries2M, L4Entries4K, U12, U21, U30};

    pub trait PageSize {
        type Entries;
        type Offset;
        fn decompose(addr_bits: u64) -> (Self::Entries, Self::Offset);
    }

    impl PageSize for L4Entries4K {
        type Entries = [u16; 4];
        type Offset = U12;

        fn decompose(addr_bits: u64) -> (Self::Entries, Self::Offset) {
            let entries = [
                ((addr_bits >> 39) & 0x1FF) as u16,
                ((addr_bits >> 30) & 0x1FF) as u16,
                ((addr_bits >> 21) & 0x1FF) as u16,
                ((addr_bits >> 12) & 0x1FF) as u16,
            ];
            let offset = U12((addr_bits & 0xFFF) as u16);
            (entries, offset)
        }
    }

    impl PageSize for L4Entries2M {
        type Entries = [u16; 3];
        type Offset = U21;

        fn decompose(addr_bits: u64) -> (Self::Entries, Self::Offset) {
            let entries = [
                ((addr_bits >> 39) & 0x1FF) as u16,
                ((addr_bits >> 30) & 0x1FF) as u16,
                ((addr_bits >> 21) & 0x1FF) as u16,
            ];
            let offset = U21((addr_bits & 0x1FFFFF) as u32);
            (entries, offset)
        }
    }

    impl PageSize for L4Entries1G {
        type Entries = [u16; 2];
        type Offset = U30;

        fn decompose(addr_bits: u64) -> (Self::Entries, Self::Offset) {
            let entries = [
                ((addr_bits >> 39) & 0x1FF) as u16,
                ((addr_bits >> 30) & 0x1FF) as u16,
            ];
            let offset = U30((addr_bits & 0x3FFFFFFF) as u32);
            (entries, offset)
        }
    }
}

impl VirtAddrExt for VirtAddr {
    fn is_canonical(&self) -> bool {
        // sign bits are bits 48-63, and they must all be the same as bit 47
        //
        // shift right 48 bits puts the 47th bit into the carry flag.
        // sign-bits + CF should be 0
        let canonical: u8;
        unsafe {
            asm!(
                "shr {bits}, 48",
                "adc {bits}, 0",
                "setz {canonical}",
                bits = inout(reg) self.0 => _,
                canonical = out(reg_byte) canonical,
            );
        }

        canonical == 1
    }

    fn page_table_index<const LEVEL: u8>(&self) -> u16 {
        assert!(LEVEL <= 4, "LEVEL must be in the range 0..=4");
        let shift = 12 + (LEVEL * 9);
        ((self.0 >> shift) & 0x1FF) as u16
    }

    fn offset_4k(&self) -> u16 {
        (self.0 & 0xFFF) as u16
    }
    fn offset_2m(&self) -> u32 {
        (self.0 & 0x1FFFFF) as u32
    }
    fn offset_1g(&self) -> u32 {
        (self.0 & 0x3FFFFFFF) as u32
    }

    fn into_parts<E: sealed::PageSize>(self) -> (E::Entries, E::Offset) {
        E::decompose(self.0)
    }
}

pub struct Context {
    pub registers: Registers,
    pub rip: u64,
    pub rflags: RFlags,
    pub rsp: u64,
    /// The code segment selector.
    pub cs: u16,
}

impl core::fmt::Display for Context {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        writeln!(f, "RIP   - {:>#018x} \tRSP - {:>#018x}", self.rip, self.rsp)?;
        writeln!(
            f,
            "FLAGS - {:>#018x} \tCS  - {:>#018x}",
            self.rflags.bits(),
            self.cs
        )?;

        writeln!(
            f,
            "RAX:   {:>#018x} \tRCX: {:>#018x}",
            self.registers.rax, self.registers.rcx
        )?;
        writeln!(
            f,
            "RDX:   {:>#018x} \tRBX: {:>#018x}",
            self.registers.rdx, self.registers.rbx
        )?;
        writeln!(
            f,
            "RSP:   {:>#018x} \tRBP: {:>#018x}",
            self.rsp, self.registers.rbp
        )?;
        writeln!(
            f,
            "RSI:   {:>#018x} \tRDI: {:>#018x}",
            self.registers.rsi, self.registers.rdi
        )?;
        writeln!(
            f,
            "R8 :   {:>#018x} \tR9 : {:>#018x}",
            self.registers.r8, self.registers.r9
        )?;
        writeln!(
            f,
            "R10:   {:>#018x} \tR11: {:>#018x}",
            self.registers.r10, self.registers.r11
        )?;
        writeln!(
            f,
            "R12:   {:>#018x} \tR13: {:>#018x}",
            self.registers.r12, self.registers.r13
        )?;
        writeln!(
            f,
            "R14:   {:>#018x} \tR15: {:>#018x}",
            self.registers.r14, self.registers.r15
        )?;

        Ok(())
    }
}