use alloc::{sync::Arc, vec::Vec};
use device_api::{device::Device, dma::DmaAllocator};
use libk::error::Error;
use libk_util::sync::spin_rwlock::IrqSafeRwLock;

use crate::device::PciDeviceInfo;

pub enum PciMatch {
    Generic(fn(&PciDeviceInfo) -> bool),
    Vendor(u16, u16),
    Class(u8, Option<u8>, Option<u8>),
}

pub struct PciDriverMatch {
    pub driver: &'static dyn PciDriver,
    pub check: PciMatch,
}

pub trait PciDriver: Sync {
    fn probe(
        &self,
        info: &PciDeviceInfo,
        dma: &Arc<dyn DmaAllocator>,
    ) -> Result<Arc<dyn Device>, Error>;
    fn driver_name(&self) -> &str;
}

impl PciMatch {
    pub fn check_device(&self, info: &PciDeviceInfo) -> bool {
        match self {
            Self::Generic(f) => f(info),
            &Self::Vendor(vendor_, device_) => {
                info.vendor_id == vendor_ && info.device_id == device_
            }
            &Self::Class(class_, Some(subclass_), Some(prog_if_)) => {
                class_ == info.class && subclass_ == info.subclass && prog_if_ == info.prog_if
            }
            &Self::Class(class_, Some(subclass_), _) => {
                class_ == info.class && subclass_ == info.subclass
            }
            &Self::Class(class_, _, _) => class_ == info.class,
        }
    }
}

pub fn register_match(pmatch: PciMatch, driver: &'static dyn PciDriver) {
    DRIVERS.write().push(PciDriverMatch {
        check: pmatch,
        driver,
    });
}

pub fn lookup_driver(info: &PciDeviceInfo) -> Option<&'static dyn PciDriver> {
    DRIVERS.read().iter().find_map(|pmatch| {
        if pmatch.check.check_device(info) {
            Some(pmatch.driver)
        } else {
            None
        }
    })
}

static DRIVERS: IrqSafeRwLock<Vec<PciDriverMatch>> = IrqSafeRwLock::new(Vec::new());
