//! x86-64 multiprocessing implementation
use core::sync::atomic::Ordering;

use ::acpi::platform::{ProcessorInfo, ProcessorState};
use kernel_arch::{Architecture, ArchitectureImpl};
use kernel_arch_x86_64::{
    mem::{
        auto_lower_address, fixed, flush_tlb_entry,
        table::{PageAttributes, PageEntry, PageTable, L1, L2},
    },
    CPU_COUNT,
};
use libk::arch::Cpu;
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress, Virtualize},
    phys,
    pointer::PhysicalRefMut,
    TableAllocatorImpl,
};
use ygg_driver_acpi::AcpiAllocator;

use crate::arch::x86_64::boot::ap_entry;

static AP_BOOTSTRAP_BIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/__x86_64_ap_boot.bin"));

const AP_STACK_PAGES: usize = 8;
const AP_BOOTSTRAP_DATA: PhysicalAddress = PhysicalAddress::from_u64(0x6000);
const AP_BOOTSTRAP_CODE: PhysicalAddress = PhysicalAddress::from_u64(0x7000);

#[repr(C)]
#[allow(dead_code)]
struct ApBootstrapData {
    cr3: PhysicalAddress,
    stack_base: usize,
    stack_size: usize,
    entry: usize,
}

unsafe fn start_ap_core(apic_id: u32) {
    assert!(ArchitectureImpl::interrupt_mask());

    let bsp_cpu = Cpu::local();
    let bsp_apic = bsp_cpu.local_apic();

    let cr3 = PhysicalAddress::from_usize(auto_lower_address(&raw const fixed::KERNEL_PML4)); // KERNEL_TABLES.lock().as_physical_address();
    let stack_base = phys::alloc_pages_contiguous(AP_STACK_PAGES)
        .unwrap()
        .virtualize();
    let stack_size = AP_STACK_PAGES * 0x1000;

    let data = ApBootstrapData {
        cr3,
        stack_base,
        stack_size,
        entry: ap_entry as usize,
    };

    let mut data_ref = PhysicalRefMut::<ApBootstrapData>::map(AP_BOOTSTRAP_DATA);
    *data_ref = data;

    let cpu_count = CPU_COUNT.load(Ordering::Acquire);

    // Send an IPI to wake up the AP
    core::arch::asm!("wbinvd");
    bsp_apic.wakeup_cpu(apic_id, AP_BOOTSTRAP_CODE);

    while cpu_count == CPU_COUNT.load(Ordering::Acquire) {
        core::hint::spin_loop();
    }

    log::info!("cpu{} up", cpu_count);
}

/// Starts up application processors specified by ACPI MADT.
///
/// # Safety
///
/// Only meant to be called once by the BSP.
pub unsafe fn start_ap_cores(info: &ProcessorInfo<AcpiAllocator>) {
    let aps = &info.application_processors;

    if aps.is_empty() {
        return;
    }

    // Temporarily identity-map the lowest 2MiB
    let mut identity_l1 = PageTable::<L1>::new_zeroed::<TableAllocatorImpl>().unwrap();
    let mut identity_l2 = PageTable::<L2>::new_zeroed::<TableAllocatorImpl>().unwrap();

    identity_l1[0] =
        PageEntry::<L1>::table(identity_l2.as_physical_address(), PageAttributes::WRITABLE);
    identity_l2[0] = PageEntry::<L2>::block(PhysicalAddress::ZERO, PageAttributes::WRITABLE);

    fixed::KERNEL_PML4[0] =
        PageEntry::table(identity_l1.as_physical_address(), PageAttributes::WRITABLE);

    // Load AP_BOOTSTRAP_CODE
    let mut code_ref = PhysicalRefMut::map_slice(AP_BOOTSTRAP_CODE, AP_BOOTSTRAP_BIN.len());
    code_ref.copy_from_slice(AP_BOOTSTRAP_BIN);

    for ap in aps.iter() {
        if ap.is_ap && ap.state == ProcessorState::WaitingForSipi {
            start_ap_core(ap.local_apic_id);
        }
    }

    // Remove the identity-map
    identity_l2[0] = PageEntry::INVALID;
    flush_tlb_entry(0);

    PageTable::free::<TableAllocatorImpl>(identity_l1);
    PageTable::free::<TableAllocatorImpl>(identity_l2);
}
