use core::ops::Range;

use alloc::{sync::Arc, vec::Vec};
use device_api::{
    device::Device,
    interrupt::{
        ExternalInterruptController, InterruptAffinity, InterruptHandler, Irq, IrqOptions,
        MessageInterruptController, MsiInfo,
    },
};
use libk::device::external_interrupt_controller;
use libk_util::{sync::spin_rwlock::IrqSafeRwLock, OneTimeInit};
use yggdrasil_abi::error::Error;

use crate::{
    capability::{MsiCapability, MsiXCapability, MsiXVectorTable},
    driver::PciDriver,
    PciAddress, PciCommandRegister, PciConfigSpace, PciConfigurationSpace, PciSegmentInfo,
};

/// Describes a PCI device
#[derive(Clone)]
pub struct PciDeviceInfo {
    /// Address of the device
    pub address: PciAddress,
    /// Class field of the configuration space
    pub class: u8,
    /// Subclass field of the configuration space
    pub subclass: u8,
    /// Prog IF field of the configuration space
    pub prog_if: u8,
    /// Vendor ID field of the configuration space
    pub vendor_id: u16,
    /// Device ID field of the configuration space
    pub device_id: u16,
    /// Configuration space access method
    pub config_space: PciConfigSpace,
    /// Describes the PCI segment this device is a part of
    pub segment: Arc<PciSegmentInfo>,

    pub(crate) interrupt_config: Arc<OneTimeInit<IrqSafeRwLock<InterruptConfig>>>,
}

pub struct InterruptConfig {
    #[allow(unused)]
    preferred_mode: PreferredInterruptMode,
    configured_mode: ConfiguredInterruptMode,
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub enum PciInterruptPin {
    A,
    B,
    C,
    D,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum PreferredInterruptMode {
    Msi(bool),
    Legacy,
}

enum ConfiguredInterruptMode {
    MsiX(
        Arc<dyn MessageInterruptController>,
        MsiXVectorTable<'static>,
    ),
    Msi(Arc<dyn MessageInterruptController>),
    LegacyPin(Arc<dyn ExternalInterruptController>, PciInterruptPin),
    #[cfg_attr(not(target_arch = "x86"), allow(unused))]
    LegacyLine(Arc<dyn ExternalInterruptController>, u8),
    None,
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub struct PciInterrupt {
    pub address: PciAddress,
    pub pin: PciInterruptPin,
}

#[derive(Clone, Copy, Debug)]
pub struct PciInterruptRoute {
    pub number: u32,
    pub options: IrqOptions,
}

#[derive(Clone)]
pub struct PciMsiRoute {
    // TODO `msi-base`
    pub controller: Arc<dyn MessageInterruptController>,
}

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum PciDeviceState {
    None,
    Probed,
    Initialized,
    Failed,
}

/// Used to store PCI bus devices which were enumerated by the kernel
pub struct PciBusDevice {
    pub(crate) info: PciDeviceInfo,
    pub(crate) device: Option<Arc<dyn Device>>,
    pub(crate) driver: Option<&'static dyn PciDriver>,
    pub(crate) state: PciDeviceState,
}

impl PciDeviceInfo {
    pub fn set_command(
        &self,
        enable_irq: bool,
        enable_mem: bool,
        enable_io: bool,
        enable_bus_master: bool,
    ) {
        let command = PciCommandRegister::from_bits_retain(self.config_space.command());
        let mut new = command;
        if enable_irq {
            new &= !PciCommandRegister::DISABLE_INTERRUPTS;
        } else {
            new |= PciCommandRegister::DISABLE_INTERRUPTS;
        }
        if enable_mem {
            new |= PciCommandRegister::ENABLE_MEMORY;
        } else {
            new &= !PciCommandRegister::ENABLE_MEMORY;
        }
        if enable_io {
            new |= PciCommandRegister::ENABLE_IO;
        } else {
            new &= !PciCommandRegister::ENABLE_IO;
        }
        if enable_bus_master {
            new |= PciCommandRegister::BUS_MASTER;
        } else {
            new &= !PciCommandRegister::BUS_MASTER;
        }
        if new != command {
            self.config_space.set_command(new.bits());
        }
    }

    pub fn init_interrupts(&self, preferred_mode: PreferredInterruptMode) -> Result<(), Error> {
        self.interrupt_config
            .try_init_with(|| {
                let configured_mode = if let PreferredInterruptMode::Msi(want_msix) = preferred_mode
                    && let Some(msi_route) = self.segment.msi_translation_map.map_msi(self.address)
                {
                    // Try to setup MSI (or MSI-x, if requested)
                    let mut result = None;
                    if want_msix
                        && let Some(mut msix) = self.config_space.capability::<MsiXCapability>()
                    {
                        if let Ok(mut vt) = msix.vector_table() {
                            if let Some(mut msi) = self.config_space.capability::<MsiCapability>() {
                                msi.set_enabled(false);
                            }

                            vt.mask_all();

                            msix.set_function_mask(false);
                            msix.set_enabled(true);

                            result = Some(ConfiguredInterruptMode::MsiX(
                                msi_route.controller.clone(),
                                vt,
                            ));
                        }
                    }

                    // Fall back to MSI if MSI-x is not available or not requested
                    if result.is_none() && self.config_space.capability::<MsiCapability>().is_some()
                    {
                        result = Some(ConfiguredInterruptMode::Msi(msi_route.controller));
                    }

                    // Fall back to legacy IRQ if nothing else works
                    if let Some(result) = result {
                        result
                    } else {
                        self.legacy_interrupt_mode()
                    }
                } else {
                    // MSI not requested or segment does not have MSI functionality
                    self.legacy_interrupt_mode()
                };

                IrqSafeRwLock::new(InterruptConfig {
                    preferred_mode,
                    configured_mode,
                })
            })
            .expect("Possible bug: double-initialization of PCI(e) interrupt config");

        Ok(())
    }

    fn legacy_interrupt_mode(&self) -> ConfiguredInterruptMode {
        let Ok(intc) = external_interrupt_controller() else {
            return ConfiguredInterruptMode::None;
        };

        // TODO this should be retrieved from interrupt map
        #[cfg(any(target_arch = "x86", rust_analyzer))]
        {
            if let Some(irq) = self.config_space.interrupt_line() {
                return ConfiguredInterruptMode::LegacyLine(intc.clone(), irq);
            }
        }

        match self.config_space.interrupt_pin() {
            Some(pin) => ConfiguredInterruptMode::LegacyPin(intc.clone(), pin),
            None => ConfiguredInterruptMode::None,
        }
    }

    pub fn map_interrupt(
        &self,
        affinity: InterruptAffinity,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<Option<MsiInfo>, Error> {
        let mut irq = self.interrupt_config.get().write();

        match &mut irq.configured_mode {
            ConfiguredInterruptMode::Msi(controller) => {
                let mut msi = self
                    .config_space
                    .capability::<MsiCapability>()
                    .ok_or(Error::InvalidOperation)?;

                let info = msi.register(controller, affinity, handler)?;

                Ok(Some(info))
            }
            ConfiguredInterruptMode::MsiX(controller, msix) => {
                let info = msix.register_range(0, 1, controller, affinity, handler)?;
                Ok(Some(info[0]))
            }
            ConfiguredInterruptMode::LegacyPin(intc, pin) => {
                self.try_map_legacy(intc.as_ref(), *pin, handler)?;
                Ok(None)
            }
            ConfiguredInterruptMode::LegacyLine(intc, irq) => {
                self.try_map_legacy_line(intc.as_ref(), *irq, handler)?;
                Ok(None)
            }
            ConfiguredInterruptMode::None => Err(Error::InvalidOperation),
        }
    }

    pub fn map_interrupt_multiple(
        &self,
        vector_range: Range<usize>,
        affinity: InterruptAffinity,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<Vec<MsiInfo>, Error> {
        let mut irq = self.interrupt_config.get().write();
        let start = vector_range.start;
        let end = vector_range.end;

        match &mut irq.configured_mode {
            ConfiguredInterruptMode::MsiX(controller, msix) => {
                msix.register_range(start, end, controller, affinity, handler)
            }
            _ => Err(Error::InvalidOperation),
        }
    }

    fn try_map_legacy(
        &self,
        intc: &dyn ExternalInterruptController,
        pin: PciInterruptPin,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<(), Error> {
        let src = PciInterrupt {
            address: self.address,
            pin,
        };
        let route = self
            .segment
            .irq_translation_map
            .map_interrupt(&src)
            .inspect_err(|e| log::warn!("Could not map PCI IRQ {pin:?}: {e:?}"))?;

        log::debug!(
            "PCI {} pin {:?} -> system IRQ #{}",
            src.address,
            src.pin,
            route.number
        );

        let irq = Irq::External(route.number);
        intc.register_irq(irq, route.options, handler)?;
        intc.enable_irq(irq)
    }

    fn try_map_legacy_line(
        &self,
        intc: &dyn ExternalInterruptController,
        line: u8,
        handler: Arc<dyn InterruptHandler>,
    ) -> Result<(), Error> {
        log::debug!("PCI {} -> IRQ#{}", self.address, line);

        let irq = Irq::External(line as u32);
        intc.register_irq(irq, Default::default(), handler)?;
        intc.enable_irq(irq)
    }
}

impl TryFrom<u32> for PciInterruptPin {
    type Error = ();

    fn try_from(value: u32) -> Result<Self, Self::Error> {
        match value {
            1 => Ok(Self::A),
            2 => Ok(Self::B),
            3 => Ok(Self::C),
            4 => Ok(Self::D),
            _ => Err(()),
        }
    }
}
