//! Simultaneous multiprocessing support for aarch64
use core::sync::atomic::Ordering;

use abi::error::Error;
use device_api::CpuBringupDevice;
use device_tree::dt::{DevTreeIndexNodePropGet, DeviceTree};
use kernel_arch_aarch64::CPU_COUNT;

use crate::arch::PLATFORM;
use crate::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,
}

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.node_by_path("/cpus").unwrap();

    cpus.children().filter_map(|cpu_node| {
        let compatible = cpu_node.prop("compatible")?;
        let id = cpu_node.prop("reg")?;
        let enable_method_str: &str = cpu_node.prop("enable-method")?;
        let enable_method = match enable_method_str {
            "psci" => CpuEnableMethod::Psci,
            _ => {
                log::warn!(
                    "Cannot enable cpu #{}: enable-method={:?} unsupported",
                    id,
                    enable_method_str
                );
                return None;
            }
        };

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

impl CpuEnableMethod {
    unsafe fn start_cpu(&self, id: usize, ip: usize, sp: usize) -> Result<(), Error> {
        match self {
            Self::Psci => {
                let psci = PLATFORM.psci.try_get().ok_or_else(|| {
                    log::warn!(
                        "cpu{} has to be enabled through PSCI, but no PSCI found",
                        id
                    );
                    Error::InvalidArgument
                })?;

                psci.start_cpu(id, ip, sp)
            }
        }
    }
}

/// 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::debug!(
            "cpu{}: enable-method={:?}, compatible={:?}",
            cpu.id,
            cpu.enable_method,
            cpu.compatible
        );

        // 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;
        // let sp = stack_pages.add(AP_STACK_PAGES * 0x1000);
        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(())
}