#[allow(dead_code)]
#[repr(C, packed)]
pub struct Entry {
    pub limit_lo: u16,
    pub base_lo: u16,
    pub base_mi: u8,
    pub access: u8,
    pub flags: u8,
    pub base_hi: u8,
}

#[allow(dead_code)]
#[repr(C, packed)]
pub struct Pointer {
    pub limit: u16,
    pub offset: usize,
}

impl Entry {
    const ACC_PRESENT: u8 = 1 << 7;
    const ACC_SYSTEM: u8 = 1 << 4;
    const ACC_EXECUTE: u8 = 1 << 3;
    const ACC_WRITE: u8 = 1 << 1;

    #[allow(unused)]
    const ACC_RING3: u8 = 3 << 5;
    #[allow(unused)]
    const ACC_ACCESS: u8 = 1 << 0;

    pub const NULL: Self = Self {
        base_lo: 0,
        base_mi: 0,
        base_hi: 0,
        access: 0,
        flags: 0,
        limit_lo: 0,
    };

    const fn new(base: u32, limit: u32, flags: u8, access: u8) -> Self {
        Self {
            base_lo: (base & 0xFFFF) as u16,
            base_mi: ((base >> 16) & 0xFF) as u8,
            base_hi: ((base >> 24) & 0xFF) as u8,
            access,
            flags: (flags & 0xF0) | (((limit >> 16) & 0xF) as u8),
            limit_lo: (limit & 0xFFFF) as u16,
        }
    }
}

#[cfg(any(target_arch = "x86", rust_analyzer))]
mod imp {
    use super::Entry;

    impl Entry {
        pub const FLAG_4K: u8 = 1 << 7;
        pub const FLAG_32: u8 = 1 << 6;

        pub const RING0_CS32: Entry = Entry::new(
            0,
            0xFFFFF,
            Entry::FLAG_4K | Entry::FLAG_32,
            Entry::ACC_ACCESS | Entry::ACC_SYSTEM | Entry::ACC_EXECUTE | Entry::ACC_PRESENT,
        );
        pub const RING0_DS32: Entry = Entry::new(
            0,
            0xFFFFF,
            Entry::FLAG_4K | Entry::FLAG_32,
            Entry::ACC_ACCESS | Entry::ACC_SYSTEM | Entry::ACC_WRITE | Entry::ACC_PRESENT,
        );
        pub const RING3_CS32: Entry = Entry::new(
            0,
            0xFFFFF,
            Entry::FLAG_4K | Entry::FLAG_32,
            Entry::ACC_SYSTEM | Entry::ACC_EXECUTE | Entry::ACC_PRESENT | Entry::ACC_RING3,
        );
        pub const RING3_DS32: Entry = Entry::new(
            0,
            0xFFFFF,
            Entry::FLAG_4K | Entry::FLAG_32,
            Entry::ACC_SYSTEM | Entry::ACC_WRITE | Entry::ACC_PRESENT | Entry::ACC_RING3,
        );
        pub const RING3_GS32: Entry = Entry::new(
            0,
            0xF,
            Entry::FLAG_32,
            Entry::ACC_SYSTEM | Entry::ACC_WRITE | Entry::ACC_PRESENT | Entry::ACC_RING3,
        );

        pub const fn tss(base: u32, limit: u32) -> Self {
            Self::new(
                base,
                limit,
                0,
                Entry::ACC_ACCESS | Entry::ACC_PRESENT | Entry::ACC_EXECUTE,
            )
        }

        /// Updates the base of this descriptor.
        ///
        /// # Safety
        ///
        /// The caller must ensure this code cannot be preempted and the segment will be properly
        /// reloaded afterwards to avoid desync between the GDT and the segment shadow.
        pub unsafe fn set_base(&mut self, value: usize) {
            use core::sync::atomic::Ordering;

            self.base_lo = (value & 0xFFFF) as u16;
            self.base_mi = ((value >> 16) & 0xFF) as u8;
            self.base_hi = ((value >> 24) & 0xFF) as u8;

            core::sync::atomic::fence(Ordering::Release);
        }
    }
}

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
mod imp {
    use alloc::boxed::Box;

    use super::{Entry, Pointer};

    #[allow(dead_code)]
    #[repr(C, packed)]
    pub struct Tss {
        _0: u32,
        rsp0: u64,
        rsp1: u64,
        rsp2: u64,
        _1: u32,
        ist1: u64,
        ist2: u64,
        ist3: u64,
        ist4: u64,
        ist5: u64,
        ist6: u64,
        ist7: u64,
        _2: u64,
        _3: u16,
        iopb_base: u16,
    }

    impl Tss {
        const NULL: Self = Self {
            _0: 0,
            rsp0: 0,
            rsp1: 0,
            rsp2: 0,
            _1: 0,
            ist1: 0,
            ist2: 0,
            ist3: 0,
            ist4: 0,
            ist5: 0,
            ist6: 0,
            ist7: 0,
            _2: 0,
            _3: 0,
            iopb_base: size_of::<Tss>() as u16,
        };
    }

    impl Entry {
        const FLAG_LONG: u8 = 1 << 5;

        const RING0_CS64: Self = Entry::new(
            0,
            0,
            Entry::FLAG_LONG,
            Entry::ACC_PRESENT | Entry::ACC_SYSTEM | Entry::ACC_EXECUTE,
        );
        const RING0_DS64: Self = Entry::new(
            0,
            0,
            0,
            Entry::ACC_PRESENT | Entry::ACC_SYSTEM | Entry::ACC_WRITE,
        );
        const RING3_DS64: Self = Entry::new(
            0,
            0,
            0,
            Entry::ACC_PRESENT | Entry::ACC_SYSTEM | Entry::ACC_RING3 | Entry::ACC_WRITE,
        );
        const RING3_CS64: Self = Entry::new(
            0,
            0,
            Entry::FLAG_LONG,
            Entry::ACC_PRESENT | Entry::ACC_SYSTEM | Entry::ACC_RING3 | Entry::ACC_EXECUTE,
        );

        const fn tss_low(base: u32, limit: u32) -> Self {
            Self::new(
                base,
                limit,
                Entry::FLAG_LONG,
                Entry::ACC_ACCESS | Entry::ACC_PRESENT | Entry::ACC_EXECUTE,
            )
        }
    }

    pub fn create_gdt() -> (&'static [Entry], &'static Tss) {
        // Won't be deallocated, so leaks are not a concern
        let tss = Box::leak(Box::new(Tss::NULL));
        let tss_addr = (tss as *mut Tss).addr();
        let mut gdt = Box::new([
            Entry::NULL,
            Entry::RING0_CS64,
            Entry::RING0_DS64,
            Entry::RING3_DS64,
            Entry::RING3_CS64,
            Entry::tss_low(tss_addr as u32, (size_of::<Tss>() - 1) as u32),
            Entry::NULL,
        ]);

        let tss_high = &mut gdt[6] as *mut _ as *mut u64;
        unsafe { tss_high.write_unaligned((tss_addr >> 32) as u64) };

        (Box::leak(gdt), tss)
    }

    pub(super) unsafe fn load_gdt(gdt: &[Entry]) {
        let gdt_addr = gdt.as_ptr().addr();
        let gdtr = Pointer {
            limit: size_of_val(gdt) as u16 - 1,
            offset: gdt_addr,
        };

        unsafe {
            core::arch::asm!(
                r#"
                wbinvd
                lgdt ({0})

                // Have to use iretq here
                mov %rsp, %rcx
                leaq 1f(%rip), %rax

                // SS:RSP
                pushq $0x10
                pushq %rcx

                // RFLAGS
                pushfq

                // CS:RIP
                pushq $0x08
                pushq %rax
                iretq
            1:
                mov $0x10, %ax
                mov %ax, %ds
                mov %ax, %es
                mov %ax, %fs
                mov %ax, %gs
                mov %ax, %ss

                mov $0x28, %ax
                ltr %ax
            "#,
                in(reg) &gdtr,
                out("rax") _,
                out("rcx") _,
                options(att_syntax)
            );
        }
    }

    /// Initializes and loads the GDT data structure for the current CPU.
    ///
    /// # Safety
    ///
    /// Intended to be called once per each CPU during their early initialization.
    pub unsafe fn init() -> usize {
        let (gdt, tss) = create_gdt();
        unsafe { load_gdt(gdt) };
        (tss as *const Tss).addr()
    }
}

#[cfg(not(target_arch = "x86"))]
pub use imp::*;