//! Simultaneous multiprocessing support for aarch64
use core::{
    ptr,
    sync::atomic::{compiler_fence, AtomicU64, Ordering},
};

use aarch64_cpu::asm::barrier;
use abi::error::Error;
use device_api::CpuBringupDevice;
use device_tree::{DeviceTree, DeviceTreeNodeExt};
use kernel_arch_aarch64::CPU_COUNT;
use libk_mm::address::{PhysicalAddress, Virtualize};

use crate::{
    arch::{aarch64::boot::SPIN_TABLE_STACK, PLATFORM},
    mem::KERNEL_VIRT_OFFSET,
};

use super::{BootStack, BOOT_STACK_SIZE};

#[link_section = ".bss"]
static AP_TRAMPOLINE_STACK: BootStack<BOOT_STACK_SIZE> = BootStack::zeroed();

#[derive(Debug)]
enum CpuEnableMethod {
    Psci,
    SpinTable(PhysicalAddress),
}

struct CpuInfo<'a> {
    id: u32,
    compatible: &'a str,
    enable_method: CpuEnableMethod,
}

fn enumerate_cpus<'a>(dt: &'a DeviceTree) -> impl Iterator<Item = CpuInfo<'a>> {
    // let cpus = dt.find_absolute("/cpus").unwrap();
    let cpus = dt.find_absolute("/cpus").unwrap();
    let address_cells = cpus.parent_address_cells();

    cpus.children().filter_map(move |cpu_node| {
        let compatible = cpu_node.prop_string("compatible")?;
        let id = cpu_node.prop_cell("reg", 1)? as u32;
        let enable_method = cpu_node.prop_string("enable-method")?;
        let enable_method = match enable_method {
            "psci" => CpuEnableMethod::Psci,
            "spin-table" => {
                let cpu_release_addr = PhysicalAddress::from_u64(
                    cpu_node.prop_cell("cpu-release-addr", address_cells)?,
                );
                CpuEnableMethod::SpinTable(cpu_release_addr)
            }
            _ => {
                log::warn!("Unknown enable method for cpu #{id}: {enable_method:?}");
                return None;
            }
        };

        Some(CpuInfo {
            id,
            compatible,
            enable_method,
        })
    })
}

impl CpuEnableMethod {
    unsafe fn start_cpu(&self, id: usize, ip: usize, sp: usize) -> Result<(), Error> {
        extern "C" {
            fn __aarch64_ap_spin_table_entry();
        }

        log::info!("Start CPU #{id}");
        match self {
            Self::Psci => {
                let psci = PLATFORM.psci.try_get().ok_or_else(|| {
                    log::warn!("cpu{id} has to be enabled through PSCI, but no PSCI found");
                    Error::DoesNotExist
                })?;

                psci.start_cpu(id, ip, sp)
            }
            &Self::SpinTable(cpu_release_addr) => {
                // Store a stack for the CPU
                SPIN_TABLE_STACK.store(sp, Ordering::Release);

                // Make sure the compiler doesn't reorder the stores
                compiler_fence(Ordering::SeqCst);
                barrier::dsb(barrier::ISHST);

                // Make the CPU jump to __aarch64_ap_spin_table_entry first
                let release_ptr =
                    ptr::with_exposed_provenance_mut::<u64>(cpu_release_addr.virtualize());
                let release_atomic = AtomicU64::from_ptr(release_ptr);
                let spin_entry_addr = __aarch64_ap_spin_table_entry as usize - KERNEL_VIRT_OFFSET;

                release_atomic.store(spin_entry_addr as u64, Ordering::Release);

                barrier::dsb(barrier::ISHST);
                barrier::isb(barrier::SY);

                // sev for implementations which spin with wfe
                aarch64_cpu::asm::sev();

                barrier::isb(barrier::SY);

                // Caller will wait for the CPU to finish its startup and proceed
                Ok(())
            }
        }
    }
}

/// Starts application processors using the method specified in the device tree.
///
/// TODO: currently does not handle systems where APs are already started before entry.
///
/// # Safety
///
/// The caller must ensure the physical memory manager was initialized, virtual memory tables are
/// set up and the function has not been called before.
pub unsafe fn start_ap_cores(dt: &DeviceTree) -> Result<(), Error> {
    extern "C" {
        fn __aarch64_ap_entry();
    }

    // Safety: safe, the stack is inside the kernel
    let sp = AP_TRAMPOLINE_STACK.top_addr() - KERNEL_VIRT_OFFSET;

    for cpu in enumerate_cpus(dt).filter(|cpu| cpu.id != 0) {
        log::info!(
            "cpu{}: enable-method={:x?}, compatible={:?}, sp={:#x}",
            cpu.id,
            cpu.enable_method,
            cpu.compatible,
            sp
        );

        // Wait for the CPU to come up
        let old_count = CPU_COUNT.load(Ordering::Acquire);

        // Safety: safe, the function is inside the kernel
        let ip = __aarch64_ap_entry as usize - KERNEL_VIRT_OFFSET;

        barrier::dsb(barrier::ISH);
        barrier::isb(barrier::SY);

        if let Err(error) = cpu.enable_method.start_cpu(cpu.id as usize, ip, sp) {
            log::error!("Couldn't start cpu{} up: {:?}", cpu.id, error);
            continue;
        }

        while CPU_COUNT.load(Ordering::Acquire) == old_count {
            aarch64_cpu::asm::wfe();
        }

        log::debug!("cpu{} is up", cpu.id);
    }

    Ok(())
}
