somehal 0.6.9

use alloc::{collections::BTreeMap, format};
use core::cell::UnsafeCell;

use aarch64_cpu::registers::ID_AA64PFR0_EL1;
use arm_gic_driver::v3::*;
use kernutil::StaticCell;
use rdrive::{PlatformDevice, module_driver, probe::OnProbeError, register::FdtInfo};

use crate::common::ioremap;

static CPU_IF: StaticCell<BTreeMap<usize, CpuInterfaceSlot>> = StaticCell::uninit();

struct CpuInterfaceSlot {
    inner: UnsafeCell<Option<CpuInterface>>,
}

// SAFETY: CPU_IF is initialized once by the BSP with all logical CPU slots
// preallocated, so the BTreeMap structure is immutable afterwards. Each CPU
// writes only its own slot during interrupt-controller initialization, and
// send_ipi reads the current CPU slot only after that CPU has initialized it.
unsafe impl Sync for CpuInterfaceSlot {}

impl CpuInterfaceSlot {
    const fn empty() -> Self {
        Self {
            inner: UnsafeCell::new(None),
        }
    }

    unsafe fn set(&self, cpu_id: usize, cpu_if: CpuInterface) {
        let slot = unsafe { &mut *self.inner.get() };
        assert!(
            slot.is_none(),
            "GICv3 CPU interface for CPU {cpu_id} is already initialized"
        );
        *slot = Some(cpu_if);
    }

    unsafe fn get(&self, cpu_id: usize) -> &CpuInterface {
        unsafe { &*self.inner.get() }
            .as_ref()
            .unwrap_or_else(|| panic!("GICv3 CPU interface for CPU {cpu_id} is not initialized"))
    }
}

pub fn with_gic(f: impl FnOnce(&mut Gic)) {
    let mut gic = super::get_gicd().lock().unwrap();
    if let Some(gic) = gic.typed_mut::<Gic>() {
        f(gic);
    }
}

module_driver!(
    name: "GICv3",
    level: ProbeLevel::PreKernel,
    priority: ProbePriority::INTC,
    probe_kinds: &[
        ProbeKind::Fdt {
            compatibles: &["arm,gic-v3"],
            on_probe: probe_gic
        }
    ],
);

fn probe_gic(info: FdtInfo<'_>, dev: PlatformDevice) -> Result<(), OnProbeError> {
    let mut reg = info.node.regs().into_iter();
    let gicd_reg = reg.next().ok_or(OnProbeError::other(format!(
        "[{}] has no reg",
        info.node.name()
    )))?;
    let gicr_reg = reg.next().unwrap();

    let gicd = ioremap(
        gicd_reg.address,
        gicd_reg.size.unwrap_or(0x1000).try_into().unwrap(),
    )
    .unwrap();
    let gicr = ioremap(
        gicr_reg.address,
        gicr_reg.size.unwrap_or(0x1000).try_into().unwrap(),
    )
    .unwrap();

    let mut gic = unsafe { Gic::new(gicd.as_ptr().into(), gicr.as_ptr().into()) };
    gic.init();
    super::set_backend(super::GicBackend::V3);

    init_cpu_interface_map();
    init_cpu_interface(&gic, 0);

    dev.register(rdif_intc::Intc::new(gic));

    Ok(())
}

/// Check if support GIC cpu interface.
pub fn is_support_icc() -> bool {
    let val = ID_AA64PFR0_EL1.get();
    // Check GIC field
    val >> 24 & 0xf > 0
}

pub fn handle_irq() -> someboot::irq::IrqId {
    let ack = ack1();

    super::_handle_irq(someboot::irq::IrqId::new(ack.to_u32() as _));

    if !ack.is_special() {
        eoi1(ack);
        if eoi_mode() {
            dir(ack);
        }
    }
    let id: u32 = ack.into();
    (id as usize).into()
}

pub fn irq_set_enable(raw: usize, enable: bool) {
    with_gic(|gic| {
        gic.set_irq_enable(unsafe { IntId::raw(raw as _) }, enable);
    });
}

pub fn send_ipi(raw: usize, target: crate::irq::IpiTarget) {
    let sgi = IntId::sgi(raw as u32);
    let target = match target {
        crate::irq::IpiTarget::Current { cpu_id: _ } => SGITarget::current(),
        crate::irq::IpiTarget::Other { cpu_id } => {
            SGITarget::list([affinity_from_mpidr(super::hardware_cpu_id(cpu_id))])
        }
        crate::irq::IpiTarget::AllExceptCurrent { .. } => SGITarget::All,
    };
    current_cpu_interface().send_sgi(sgi, target);
}

fn affinity_from_mpidr(mpidr: usize) -> Affinity {
    Affinity {
        aff0: (mpidr & 0xff) as u8,
        aff1: ((mpidr >> 8) & 0xff) as u8,
        aff2: ((mpidr >> 16) & 0xff) as u8,
        aff3: ((mpidr >> 32) & 0xff) as u8,
    }
}

pub fn init_cpu() {
    let cpu_id = someboot::smp::early_current_cpu_idx();
    with_gic(|gic| init_cpu_interface(gic, cpu_id));

    debug!("GICCv3 initialized");
}

fn init_cpu_interface_map() {
    let mut cpu_if = BTreeMap::new();
    for cpu_id in 0..someboot::smp::cpu_count() {
        cpu_if.insert(cpu_id, CpuInterfaceSlot::empty());
    }
    CPU_IF.init(cpu_if);
}

fn init_cpu_interface(gic: &Gic, cpu_id: usize) {
    let mut cpu = gic.cpu_interface();
    cpu.init_current_cpu().unwrap();
    #[cfg(feature = "hv")]
    cpu.set_eoi_mode(true);

    // SAFETY: CPU_IF was preallocated during BSP probe. Each CPU initializes
    // only its own logical CPU slot before it can send SGIs through that slot.
    unsafe { cpu_interface_slot(cpu_id).set(cpu_id, cpu) };
}

fn current_cpu_interface() -> &'static CpuInterface {
    let cpu_id = someboot::smp::early_current_cpu_idx();
    // SAFETY: send_ipi is only valid after the current CPU has completed
    // interrupt-controller initialization and stored its CpuInterface.
    unsafe { cpu_interface_slot(cpu_id).get(cpu_id) }
}

fn cpu_interface_slot(cpu_id: usize) -> &'static CpuInterfaceSlot {
    CPU_IF
        .get(&cpu_id)
        .unwrap_or_else(|| panic!("GICv3 CPU interface slot for CPU {cpu_id} is not registered"))
}