#![no_std]

use alloc::sync::Arc;
use device_api::{device::Device, dma::DmaAllocator};
use libk::error::Error;
use rtl8139::Rtl8139;
use rtl8168::Rtl8168;
use ygg_driver_pci::{
    capability::{
        DevicePowerState, PciExpressCapability, PcieLinkControl, PowerManagementCapability,
    },
    device::{PciDeviceInfo, PreferredInterruptMode},
    macros::pci_driver,
    PciBaseAddress, PciCommandRegister, PciConfigurationSpace,
};

extern crate alloc;

pub mod rtl8139;
pub mod rtl8168;

pci_driver! {
    matches: [device (0x10EC:0x8139)],
    driver: {
        fn driver_name(&self) -> &str {
            "rtl8139"
        }

        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
            info.init_interrupts(PreferredInterruptMode::Msi(false))?;

            if let Some(mut pcie) = info.config_space.capability::<PciExpressCapability>() {
                let mut lcr = pcie.link_control();
                lcr.remove(PcieLinkControl::ASPM_MASK | PcieLinkControl::ECPM);
                pcie.set_link_control(lcr);
            }

            // Enable MMIO + interrupts + bus mastering
            let mut command = info.config_space.command();
            command |= (PciCommandRegister::BUS_MASTER | PciCommandRegister::ENABLE_MEMORY).bits();
            command &= !(PciCommandRegister::ENABLE_IO | PciCommandRegister::DISABLE_INTERRUPTS).bits();
            info.config_space.set_command(command);

            let base = info
                .config_space
                .bar(1)
                .and_then(PciBaseAddress::as_memory)
                .ok_or(Error::InvalidArgument)?;

            let device = Rtl8139::new(dma.clone(), base, info.clone())?;
            Ok(Arc::new(device))
        }
    }
}

pci_driver! {
    matches: [device (0x10EC:0x8168)],
    driver: {
        fn driver_name(&self) -> &str {
            "rtl816x"
        }

        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
            // Enable MMIO + interrupts + bus mastering
            info.set_command(false, true, true, true);

            let base = info
                .config_space
                .bar(2)
                .and_then(PciBaseAddress::as_memory)
                .ok_or(Error::InvalidArgument)?;

            if let Some(power) = info.config_space.capability::<PowerManagementCapability>() {
                power.set_device_power_state(DevicePowerState::D0);
                power.set_pme_en(true);
            }

            if let Some(mut pcie) = info.config_space.capability::<PciExpressCapability>() {
                let mut lcr = pcie.link_control();
                lcr.remove(PcieLinkControl::ASPM_MASK | PcieLinkControl::ECPM);
                pcie.set_link_control(lcr);
            }

            let device = Rtl8168::new(dma.clone(), base, info.clone())?;
            Ok(Arc::new(device))
        }
    }
}
