//! RISC-V SMP implementation

use core::{mem::MaybeUninit, sync::atomic::Ordering};

use abi::error::Error;
use device_api::interrupt::{IpiDeliveryTarget, IpiMessage};
use device_tree::{DeviceTree, DeviceTreeNodeExt};
use kernel_arch_riscv64::{
    boot_hart_id, mem::auto_lower_address, registers::SIP, sbi, ArchitectureImpl, CPU_COUNT,
};
use libk::arch::Cpu;
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    PageBox,
};
use tock_registers::interfaces::ReadWriteable;

use crate::panic;

#[allow(missing_docs)]
pub const SECONDARY_STACK_SIZE: usize = 32768;

#[repr(C)]
#[allow(missing_docs)]
pub struct SecondaryContext {
    // 0x00
    pub stack_top: PhysicalAddress,

    pub stack: PageBox<[MaybeUninit<u8>]>,
    pub hart_id: u64,
}

fn start_secondary_hart(hart_id: u64) -> Result<(), Error> {
    extern "C" {
        fn __riscv64_ap_entry();
    }

    let start_addr = auto_lower_address(__riscv64_ap_entry as *const fn());
    let stack = PageBox::<u8>::new_uninit_slice(SECONDARY_STACK_SIZE)?;
    let stack_top = unsafe { PageBox::as_physical_address(&stack).add(SECONDARY_STACK_SIZE) };
    let context = PageBox::new(SecondaryContext {
        stack,
        stack_top,
        hart_id,
    })?;
    let a1 = context.into_physical_raw();

    log::info!(
        "Start secondary hart {hart_id}: pc={:#x}, a1={:#x}",
        start_addr,
        a1,
    );

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

    sbi::sbi_hart_start(hart_id, start_addr as u64, a1.into_u64())?;

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

    Ok(())
}

/// Main function to setup all secondary HARTs
pub fn start_secondary_harts(dt: &DeviceTree) -> Result<(), Error> {
    log::info!("Setting up secondary harts");

    let boot_hart_id = boot_hart_id() as u32;
    let cpus = dt.find_absolute("/cpus").ok_or(Error::DoesNotExist)?;

    // Find the boot hart
    let boot_hart_isa = cpus
        .children()
        .find(|child| child.prop_cell_usize("reg") == Some(boot_hart_id as usize))
        .and_then(|child| child.prop_string("riscv,isa"))
        .ok_or(Error::DoesNotExist)
        .inspect_err(|_| log::error!("Could not find the boot hart {boot_hart_id} in the dtb"))?;

    // Print a list of HARTs
    for cpu in cpus.children() {
        let Some(reg) = cpu.prop_cell_usize("reg") else {
            continue;
        };
        let isa = cpu.prop_string("riscv,isa").unwrap_or("???");

        let letter = if reg == boot_hart_id as usize {
            '@'
        } else if isa == boot_hart_isa {
            '+'
        } else {
            '-'
        };

        log::info!("{letter} {reg}: {isa}");
    }

    for cpu in cpus.children() {
        let Some(reg) = cpu.prop_cell_usize("reg") else {
            continue;
        };
        let isa = cpu.prop_string("riscv,isa").unwrap_or("???");

        if reg == boot_hart_id as usize || isa != boot_hart_isa {
            continue;
        }

        if let Err(error) = start_secondary_hart(reg as u64) {
            log::error!("hart {reg} start error: {error:?}");
        }
    }

    // Can get rid of lower half now
    // unsafe { mem::unmap_lower_half() };

    Ok(())
}

/// Sends an IPI to the target HART(s)
pub fn send_ipi(target: IpiDeliveryTarget, msg: IpiMessage) -> Result<(), Error> {
    let local = Cpu::local();

    let mut hart_mask = 0;
    match target {
        IpiDeliveryTarget::Specific(_) => todo!(),
        IpiDeliveryTarget::ThisCpu => todo!(),
        IpiDeliveryTarget::OtherCpus => {
            ArchitectureImpl::for_each_hart(|hart_id, _, queue| {
                if hart_id != local.id() {
                    hart_mask |= 1 << hart_id;
                    queue.push(msg);
                }
            });
        }
    }

    log::debug!("send_ipi({hart_mask:#x}, {msg:?})");
    sbi::sbi_send_ipi(hart_mask, 0)
}

/// Handles received IPIs
pub fn handle_ipi() {
    let local = Cpu::local();
    if let Some(message) = local.get_ipi() {
        match message {
            IpiMessage::Panic => panic::panic_secondary(),
            IpiMessage::Shutdown => todo!(),
        }
    } else {
        log::warn!("Spurious IPI received");
    }
    SIP.modify(SIP::SSIP::CLEAR);
}
