//! ARM GICv2m extension support for MSI/MSI-x
use core::{mem::offset_of, ops::Range};

use abi::error::Error;
use alloc::{sync::Arc, vec::Vec};
use device_api::{
    device::{Device, DeviceInitContext},
    interrupt::{
        ExternalInterruptController, InterruptAffinity, InterruptHandler, Irq, IrqLevel,
        IrqOptions, IrqTrigger, IrqVector, MessageInterruptController, MsiInfo,
    },
};
use device_tree::driver::{device_tree_driver, Node, ProbeContext};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use tock_registers::{
    interfaces::Readable,
    register_bitfields, register_structs,
    registers::{ReadOnly, WriteOnly},
};

use crate::arch::aarch64::gic::GIC_SPI_START;

register_bitfields! {
    u32,
    GICM_TYPER [
        VALID       OFFSET(31)  NUMBITS(1)  [],
        CLR         OFFSET(30)  NUMBITS(1)  [],
        SR          OFFSET(29)  NUMBITS(1)  [],
        INTID       OFFSET(16)  NUMBITS(13) [],
        NUMSPIS     OFFSET(0)   NUMBITS(11) [],
    ],
    GICM_IIDR [
        VARIANT     OFFSET(16)  NUMBITS(4)  [],
        REVISION    OFFSET(12)  NUMBITS(4)  [],
        IMPLEMENTER OFFSET(0)   NUMBITS(12) [],
    ],
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x0000 => _0),
        (0x0008 => GICM_TYPER: ReadOnly<u32, GICM_TYPER::Register>),
        (0x000C => _1),
        (0x0040 => GICM_SETSPI_NSR: WriteOnly<u32>),
        (0x0044 => _2),
        (0x0048 => GICM_CLRSPI_NSR: WriteOnly<u32>),
        (0x004C => _3),
        (0x0050 => GICM_SETSPI_SR: WriteOnly<u32>),
        (0x0054 => _4),
        (0x0058 => GICM_CLRSPI_SR: WriteOnly<u32>),
        (0x005C => _5),
        (0x0FCC => GICM_IIDR: ReadOnly<u32, GICM_IIDR::Register>),
        (0x0FD0 => _6),
        (0x1000 => @END),
    }
}

struct MsiEntry {
    handler: Arc<dyn InterruptHandler>,
    vector: usize,
}

struct Inner {
    irq_range: Range<u32>,
    msi_map: IrqSafeRwLock<Vec<Option<MsiEntry>>>,
}

/// GICv2m Register Frame, allows handling MSI(-x)-based interrupts via its parent GIC
pub struct Gicv2m {
    base: PhysicalAddress,
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    parent: Arc<dyn ExternalInterruptController>,

    msi_base_override: Option<u32>,
    msi_count_override: Option<usize>,

    inner: OneTimeInit<Inner>,
}

impl Device for Gicv2m {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let regs = self.regs.lock();
        log::info!("gicv2m: init @ {:#x}", self.base);

        // NOTE qemu's gicv2m TYPER.VALID = 0 for some reason, even though it contains proper
        //      values
        let quirk_qemu = (regs.GICM_IIDR.get() >> 20) == b'Q' as u32;

        // If GICM_TYPER.VALID == 0, use the base provided via device tree (or fail if neither is
        // available)
        let valid = quirk_qemu || regs.GICM_TYPER.matches_all(GICM_TYPER::VALID::SET);

        let irq_base = if valid {
            regs.GICM_TYPER.read(GICM_TYPER::INTID)
        } else {
            self.msi_base_override
                .ok_or(Error::InvalidArgument)
                .inspect_err(|_| {
                    log::error!("gicv2m: typer is not valid and no base provided in the FDT")
                })?
        };
        let irq_count = if valid {
            regs.GICM_TYPER.read(GICM_TYPER::NUMSPIS)
        } else {
            self.msi_count_override
                .ok_or(Error::InvalidArgument)
                .inspect_err(|_| {
                    log::error!("gicv2m: typer is not valid and no irq count provided in the FDT")
                })? as u32
        };

        log::info!("gicv2m: irq base {irq_base}, irq count {irq_count}");
        let irq_range = irq_base..irq_base + irq_count;
        self.inner.init(Inner {
            irq_range: irq_range.clone(),
            msi_map: IrqSafeRwLock::new((0..irq_count as usize).map(|_| None).collect()),
        });

        for irq in irq_range {
            // Map SPI in parent
            let irq = Irq::External(irq - 32);
            self.parent.register_irq(
                irq,
                IrqOptions {
                    trigger: IrqTrigger::Edge,
                    level: IrqLevel::ActiveHigh,
                },
                self.clone(),
            )?;
        }

        Ok(())
    }

    fn display_name(&self) -> &str {
        "GICv2m Frame"
    }
}

impl InterruptHandler for Gicv2m {
    fn handle_irq(self: Arc<Self>, vector: IrqVector) -> bool {
        let IrqVector::Irq(Irq::External(spi)) = vector else {
            return false;
        };
        let Some(inner) = self.inner.try_get() else {
            return false;
        };
        let msi_map = inner.msi_map.read();
        let index = (spi + GIC_SPI_START - inner.irq_range.start) as usize;
        if index >= inner.irq_range.len() {
            return false;
        }
        if let Some(entry) = &msi_map[index] {
            entry
                .handler
                .clone()
                .handle_irq(IrqVector::Msi(entry.vector))
        } else {
            false
        }
    }
}

impl MessageInterruptController for Gicv2m {
    fn handle_msi(&self, _vector: usize) {
        todo!()
    }

    fn register_msi_range(
        self: Arc<Self>,
        range: &mut [MsiInfo],
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<(), Error> {
        let inner = self.inner.get();
        let mut msi_map = inner.msi_map.write();
        // let map = self.irq_map.write();
        // let irq_range = self.irq_range.try_get().ok_or(Error::DoesNotExist)?.clone();
        let irq_base = inner.irq_range.start;

        let mut allocated_range = None;
        for irq in inner.irq_range.clone() {
            let mut taken = false;
            for i in 0..range.len() {
                let index = (irq - irq_base) as usize + i;
                if msi_map[index].is_some() {
                    taken = true;
                    break;
                }
            }
            if taken {
                continue;
            }

            allocated_range = Some(irq..irq + range.len() as u32);
            break;
        }

        let map_range = allocated_range
            .ok_or(Error::InvalidArgument)
            .inspect_err(|_| log::warn!("gicv2m: no free irq slots"))?;

        for (i, (info, irq)) in core::iter::zip(range, map_range).enumerate() {
            log::info!("gicv2m: map msi irq {irq} -> {}", handler.display_name());
            let index = (irq - irq_base) as usize;

            msi_map[index] = Some(MsiEntry {
                vector: i,
                handler: handler.clone(),
            });

            *info = MsiInfo {
                address: self
                    .base
                    .add(offset_of!(Regs, GICM_SETSPI_NSR))
                    .into_usize(),
                value: irq,
                vector: i,
                affinity: InterruptAffinity::Any,
            };
            self.parent.enable_irq(Irq::External(irq - 32))?;
        }

        Ok(())
    }
}

device_tree_driver! {
    compatible: ["arm,gic-v2m-frame"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let base = node.map_base(context, 0)?;
            let parent = node.parent()?.as_interrupt_controller()?;

            let msi_base_override = node.prop_usize("arm,msi-base-spi").map(|x| x as u32);
            let msi_count_override = node.prop_usize("arm,msi-num-spis");

            let regs = unsafe { DeviceMemoryIo::map(base, Default::default()) }.inspect_err(|e| {
                log::error!("gicv2m: could not map memory at {base:#x}: {e:?}");
            }).ok()?;

            let intc = Arc::new(Gicv2m {
                base,
                parent,
                msi_base_override,
                msi_count_override,
                regs: IrqSafeSpinlock::new(regs),
                inner: OneTimeInit::new()
            });

            node.make_msi_controller(intc.clone());
            Some(intc)
        }
    }
}
