use core::{
    alloc::Layout,
    pin::Pin,
    ptr::{addr_of, NonNull},
};

use super::{enable_sse, percpu::init_percpu_area_thiscpu, GDTEntry, InterruptControl, GDT};

#[repr(C)]
#[derive(Debug, Clone, Copy)]
#[allow(non_camel_case_types)]
struct TSS_SP {
    low: u32,
    high: u32,
}

#[repr(C)]
pub struct TSS {
    _reserved1: u32,
    rsp: [TSS_SP; 3],
    _reserved2: u32,
    _reserved3: u32,
    ist: [TSS_SP; 7],
    _reserved4: u32,
    _reserved5: u32,
    _reserved6: u16,
    iomap_base: u16,
}

impl TSS {
    pub fn new() -> Self {
        Self {
            _reserved1: 0,
            rsp: [TSS_SP { low: 0, high: 0 }; 3],
            _reserved2: 0,
            _reserved3: 0,
            ist: [TSS_SP { low: 0, high: 0 }; 7],
            _reserved4: 0,
            _reserved5: 0,
            _reserved6: 0,
            iomap_base: 0,
        }
    }

    pub fn set_rsp0(&mut self, rsp: u64) {
        self.rsp[0].low = rsp as u32;
        self.rsp[0].high = (rsp >> 32) as u32;
    }
}

/// Architecture-specific per-cpu status.
#[allow(dead_code)]
pub struct CPUStatus {
    id: usize,
    gdt: GDT,
    tss: TSS,

    percpu_area: NonNull<u8>,
    pub interrupt: InterruptControl,
}

impl CPUStatus {
    pub unsafe fn new_thiscpu<F>(allocate: F) -> Self
    where
        F: FnOnce(Layout) -> NonNull<u8>,
    {
        const PAGE_SIZE: usize = 0x1000;
        extern "C" {
            static PERCPU_PAGES: usize;
            fn _PERCPU_DATA_START();
        }

        let percpu_area = allocate(Layout::from_size_align_unchecked(
            PERCPU_PAGES * PAGE_SIZE,
            PAGE_SIZE,
        ));

        percpu_area.copy_from_nonoverlapping(
            NonNull::new(_PERCPU_DATA_START as *mut u8).unwrap(),
            PERCPU_PAGES * PAGE_SIZE,
        );

        let (interrupt_control, cpuid) = InterruptControl::new();

        init_percpu_area_thiscpu(percpu_area);
        Self {
            id: cpuid,
            gdt: GDT::new(),
            tss: TSS::new(),
            percpu_area,
            interrupt: interrupt_control,
        }
    }

    /// Load GDT and TSS in place.
    ///
    /// # Safety
    /// Make sure preemption and interrupt are disabled before calling this function.
    pub unsafe fn init(self: Pin<&mut Self>) {
        enable_sse();

        // SAFETY: We don't move the object.
        let self_mut = self.get_unchecked_mut();

        let tss_addr = addr_of!(self_mut.tss);
        self_mut.gdt.set_tss(tss_addr as u64);
        self_mut.gdt.load();

        // SAFETY: `self` is pinned, so are its fields.
        Pin::new_unchecked(&mut self_mut.interrupt).setup_idt();
        self_mut.interrupt.setup_timer();
    }

    /// Bootstrap all CPUs.
    /// This should only be called on the BSP.
    pub unsafe fn bootstrap_cpus(&self) {
        self.interrupt.send_sipi();
    }

    pub unsafe fn set_rsp0(&mut self, rsp: u64) {
        self.tss.set_rsp0(rsp);
    }

    pub unsafe fn set_tls32(&mut self, desc: GDTEntry) {
        self.gdt.set_tls32(desc);
    }

    pub fn cpuid(&self) -> usize {
        self.id
    }
}

#[macro_export]
macro_rules! define_smp_bootstrap {
    ($cpu_count:literal, $ap_entry:ident, $alloc_kstack:tt) => {
        #[no_mangle]
        static BOOT_SEMAPHORE: core::sync::atomic::AtomicU64 =
            core::sync::atomic::AtomicU64::new(0);
        #[no_mangle]
        static BOOT_STACK: core::sync::atomic::AtomicU64 =
            core::sync::atomic::AtomicU64::new(0);

        #[no_mangle]
        static CPU_COUNT: core::sync::atomic::AtomicU64 =
            core::sync::atomic::AtomicU64::new(0);

        core::arch::global_asm!(
            r#"
        .pushsection .stage1.smp
        .code16
        .globl ap_bootstrap
        .type ap_bootstrap, @function
        ap_bootstrap:
            ljmp $0x0, $.Lap1

        .Lap1:
            # we use the shared gdt for cpu bootstrapping
            lgdt .Lshared_gdt_desc

            # set msr
            mov $0xc0000080, %ecx
            rdmsr
            or $0x901, %eax # set LME, NXE, SCE
            wrmsr

            # set cr4
            mov %cr4, %eax
            or $0xa0, %eax # set PAE, PGE
            mov %eax, %cr4

            # load new page table
            mov ${KERNEL_PML4}, %eax
            mov %eax, %cr3

            mov %cr0, %eax
            // SET PE, WP, PG
            or $0x80010001, %eax
            mov %eax, %cr0

            ljmp $0x08, $.Lap_bootstrap_end

        .align 16
        .Lshared_gdt_desc:
            .8byte 0x0000000000005f

        .code64
        .Lap_bootstrap_end:
            mov $0x10, %ax
            mov %ax, %ds
            mov %ax, %es
            mov %ax, %ss

            xor %rsp, %rsp
            xor %rax, %rax
            inc %rax
        1:
            xchg %rax, {BOOT_SEMAPHORE}
            cmp $0, %rax
            je 1f
            pause
            jmp 1b

        1:
            mov {BOOT_STACK}, %rsp # Acquire
            cmp $0, %rsp
            jne 1f
            pause
            jmp 1b

        1:
            xor %rax, %rax
            mov %rax, {BOOT_STACK} # Release
            xchg %rax, {BOOT_SEMAPHORE}

            lock incq {CPU_COUNT}

            xor %rbp, %rbp
            push %rbp # NULL return address
            jmp {AP_ENTRY}
            .popsection
            "#,
            KERNEL_PML4 = const 0x2000,
            BOOT_SEMAPHORE = sym BOOT_SEMAPHORE,
            BOOT_STACK = sym BOOT_STACK,
            CPU_COUNT = sym CPU_COUNT,
            AP_ENTRY = sym $ap_entry,
            options(att_syntax),
        );

        pub unsafe fn wait_cpus_online() {
            use core::sync::atomic::Ordering;
            while CPU_COUNT.load(Ordering::Acquire) != $cpu_count - 1 {
                if BOOT_STACK.load(Ordering::Acquire) == 0 {
                    let stack_bottom = $alloc_kstack as u64;
                    BOOT_STACK.store(stack_bottom, Ordering::Release);
                }
                $crate::pause();
            }
        }
    };
}