use alloc::{sync::Arc, vec::Vec};
use device_api::{
    device::{Device, DeviceInitContext},
    interrupt::{
        ExternalInterruptController, FixedInterruptTable, InterruptHandler, InterruptTable, Irq,
        IrqHandle, IrqOptions, IrqVector,
    },
};
use device_tree::{
    driver::{
        device_tree_driver, lookup_phandle, DeviceTreeInterruptController, Node, ProbeContext,
    },
    DeviceTreePropertyRead, TProp,
};
use kernel_arch_riscv64::boot_hart_id;
use libk::{arch::Cpu, device::register_external_interrupt_controller};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::{sync::spin_rwlock::IrqSafeRwLock, OneTimeInit};
use tock_registers::{
    interfaces::{Readable, Writeable},
    register_structs,
    registers::{ReadOnly, ReadWrite},
};
use yggdrasil_abi::{error::Error, primitive_enum};

const MAX_IRQS: usize = 1024;

const ENABLE_BASE: usize = 0x2000;
const ENABLE_STRIDE: usize = 0x80;
const CONTROL_BASE: usize = 0x200000;
const CONTROL_STRIDE: usize = 0x1000;

primitive_enum! {
    enum ContextMode: u32 {
        SoftU = 0,
        SoftS = 1,
        SoftH = 2,
        SoftM = 3,
        TimerU = 4,
        TimerS = 5,
        TimerH = 6,
        TimerM = 7,
        ExternalU = 8,
        ExternalS = 9,
        ExternalH = 10,
        ExternalM = 11,
    }
}

register_structs! {
    #[allow(non_snake_case)]
    CommonRegs {
        (0x0000 => _0),
        (0x0004 => PRIORITY: [ReadWrite<u32>; 1023]),
        (0x1000 => PENDING: [ReadOnly<u32>; 32]),
        (0x1080 => @END),
    }
}

register_structs! {
    #[allow(non_snake_case)]
    ContextEnableRegs {
        (0x0000 => ENABLE: [ReadWrite<u32>; 32]),
        (0x0080 => @END),
    }
}

register_structs! {
    #[allow(non_snake_case)]
    ContextControlRegs {
        (0x0000 => THRESHOLD: ReadWrite<u32>),
        (0x0004 => CLAIM: ReadWrite<u32>),
        (0x0008 => @END),
    }
}

struct Context {
    enable: IrqSafeRwLock<DeviceMemoryIo<'static, ContextEnableRegs>>,
    control: IrqSafeRwLock<DeviceMemoryIo<'static, ContextControlRegs>>,
    // TODO scale the table depending on effective MAX_IRQS value
    table: IrqSafeRwLock<FixedInterruptTable<128>>,
}

struct Inner {
    #[allow(unused)]
    common: IrqSafeRwLock<DeviceMemoryIo<'static, CommonRegs>>,
}

struct HartContext {
    hart: u32,
    index: usize,
    context: OneTimeInit<Context>,
}

/// RISC-V Platform-Level Interrupt Controller (PLIC) device
pub struct Plic {
    base: PhysicalAddress,
    max_irqs: usize,
    // hart id -> context map
    context_map: Vec<HartContext>,
    inner: OneTimeInit<Inner>,
}

impl Plic {
    fn hart_context(&self, hart: u32) -> Option<&HartContext> {
        self.context_map.iter().find(|c| c.hart == hart)
    }

    fn validate_irq(&self, irq: Irq) -> Result<u32, Error> {
        let Irq::External(irq) = irq else {
            log::error!("plic: irq {irq:?} is not an external interrupt");
            return Err(Error::InvalidArgument);
        };
        if irq == 0 {
            log::error!("plic: irq cannot be zero");
            return Err(Error::InvalidArgument);
        }
        if irq as usize >= self.max_irqs {
            log::error!("plic: irq ({}) >= max_irqs ({})", irq, self.max_irqs);
            return Err(Error::InvalidArgument);
        }
        Ok(irq)
    }
}

impl ContextEnableRegs {
    fn enable_irq(&self, irq: u32) {
        let reg = &self.ENABLE[irq as usize / 32];
        reg.set(reg.get() | (1 << (irq % 32)));
    }
}

impl ExternalInterruptController for Plic {
    fn enable_irq(&self, irq: Irq) -> Result<(), Error> {
        // TODO balance IRQs between harts?
        let irq = self.validate_irq(irq)?;
        let bsp_hart_id = boot_hart_id() as u32;
        let context = self
            .hart_context(bsp_hart_id)
            .ok_or(Error::InvalidArgument)
            .inspect_err(|_| log::error!("plic: no context for hart {bsp_hart_id}"))?
            .context
            .get();

        let enable = context.enable.write();

        enable.enable_irq(irq);

        Ok(())
    }

    fn register_irq(
        &self,
        irq: Irq,
        _options: IrqOptions,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<(), Error> {
        let bsp_hart_id = boot_hart_id() as u32;
        let irq = self.validate_irq(irq)?;
        let context = self
            .hart_context(bsp_hart_id)
            .ok_or(Error::InvalidArgument)
            .inspect_err(|_| log::error!("plic: no context for hart {bsp_hart_id}"))?
            .context
            .get();
        let mut table = context.table.write();
        log::info!(
            "Bind irq #{irq} -> hart {bsp_hart_id}, {:?}",
            handler.display_name()
        );

        table.insert(irq as usize, handler)?;

        Ok(())
    }

    fn handle_pending_irqs(&self) {
        let hart_id = Cpu::local().id();
        let Some(context) = self.hart_context(hart_id) else {
            log::warn!("plic: irq on hart without a context: {hart_id}");
            return;
        };
        let context = context.context.get();

        let control = context.control.write();
        let table = context.table.read();

        loop {
            let irq = control.CLAIM.get();
            if irq == 0 {
                break;
            }
            let vector = IrqVector::Irq(Irq::External(irq));

            if let Some(handler) = table.handler(irq as usize) {
                handler.clone().handle_irq(vector);
            } else {
                log::warn!("plic: no handler for IRQ #{irq}");
            }

            // Done servicing
            control.CLAIM.set(irq);
        }
    }
}

impl Device for Plic {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        log::info!("Initialize RISC-V PLIC");

        let common = DeviceMemoryIo::<CommonRegs>::map(self.base, Default::default())?;

        for i in 0..self.max_irqs - 1 {
            common.PRIORITY[i].set(3);
        }

        for context in self.context_map.iter() {
            let enable_offset = ENABLE_BASE + context.index * ENABLE_STRIDE;
            let control_offset = CONTROL_BASE + context.index * CONTROL_STRIDE;

            log::info!(
                "* HART {}: context {}, enable={:#x}, control={:#x}",
                context.hart,
                context.index,
                enable_offset,
                control_offset
            );
            let enable = DeviceMemoryIo::<ContextEnableRegs>::map(
                self.base.add(enable_offset),
                Default::default(),
            )?;
            let control = DeviceMemoryIo::<ContextControlRegs>::map(
                self.base.add(control_offset),
                Default::default(),
            )?;

            for i in 0..self.max_irqs.div_ceil(32) {
                enable.ENABLE[i].set(0);
            }
            control.THRESHOLD.set(0);

            context.context.init(Context {
                enable: IrqSafeRwLock::new(enable),
                control: IrqSafeRwLock::new(control),
                table: IrqSafeRwLock::new(FixedInterruptTable::new()),
            });
        }

        self.inner.init(Inner {
            common: IrqSafeRwLock::new(common),
        });

        register_external_interrupt_controller(self);

        Ok(())
    }

    fn display_name(&self) -> &str {
        "RISC-V PLIC"
    }
}

impl DeviceTreeInterruptController for Plic {
    fn map_interrupt(self: Arc<Self>, property: &TProp, offset: usize) -> Option<IrqHandle> {
        let num = property.read_cell(offset, 1)?;
        Some(IrqHandle {
            irq: Irq::External(num as _),
            options: IrqOptions::default(),
            intc: self.clone(),
        })
    }

    fn as_interrupt_controller(self: Arc<Self>) -> Arc<dyn ExternalInterruptController> {
        self
    }
}

fn map_context_to_hart(target: u32) -> Option<u32> {
    let hart_intc = lookup_phandle(target, false)?;
    let parent = hart_intc.parent()?;
    let reg = parent.prop_usize("reg")?;
    Some(reg as u32)
}

device_tree_driver! {
    compatible: ["starfive,jh7110-plic", "sifive,plic-1.0.0", "riscv,plic0"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let base = node.map_base(context, 0)?;
            let ndev = node.prop_usize("riscv,ndev")?;
            let iext = node.property("interrupts-extended")?;

            let max_irqs = MAX_IRQS.min(ndev);

            // Parse the context -> hart mapping, only select S-mode external interrupts
            let mut context_map = Vec::new();
            let mut mapped_harts = 0u64;
            for (context, (target, mode)) in iext.iter_cells((1, 1)).enumerate() {
                let (Ok(mode), Some(hart_id)) = (
                    ContextMode::try_from(mode as u32),
                    map_context_to_hart(target as u32)
                ) else {
                    continue;
                };
                if mode != ContextMode::ExternalS {
                    continue;
                }
                // Don't map the same hart to two contexts
                if mapped_harts & (1 << hart_id) == 0 {
                    mapped_harts |= 1 << hart_id;
                    context_map.push(HartContext {
                        hart: hart_id,
                        index: context,
                        context: OneTimeInit::new()
                    });
                }
            }

            if context_map.is_empty() {
                log::warn!("{:?}: could not map any contexts to harts", node.name());
                return None;
            }

            let intc = Arc::new(Plic {
                base,
                max_irqs,
                context_map,
                inner: OneTimeInit::new(),
            });

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