//! Defines types and probe methods of all supported devices.

#![allow(unused_imports, dead_code)]

use ax_driver_base::DeviceType;
#[cfg(feature = "bus-pci")]
use ax_driver_pci::{ConfigurationAccess, DeviceFunction, DeviceFunctionInfo, PciRoot};

pub use super::dummy::*;
use crate::AxDeviceEnum;
#[cfg(feature = "virtio")]
use crate::virtio::{self, VirtIoDevMeta};

pub trait DriverProbe {
    fn probe_global() -> Option<AxDeviceEnum> {
        None
    }

    #[cfg(bus = "mmio")]
    fn probe_mmio(_mmio_base: usize, _mmio_size: usize) -> Option<AxDeviceEnum> {
        None
    }

    #[cfg(bus = "pci")]
    fn probe_pci<C: ConfigurationAccess>(
        _root: &mut PciRoot<C>,
        _bdf: DeviceFunction,
        _dev_info: &DeviceFunctionInfo,
    ) -> Option<AxDeviceEnum> {
        None
    }
}

#[cfg(net_dev = "virtio-net")]
register_net_driver!(
    <virtio::VirtIoNet as VirtIoDevMeta>::Driver,
    <virtio::VirtIoNet as VirtIoDevMeta>::Device
);

#[cfg(block_dev = "virtio-blk")]
register_block_driver!(
    <virtio::VirtIoBlk as VirtIoDevMeta>::Driver,
    <virtio::VirtIoBlk as VirtIoDevMeta>::Device
);

#[cfg(display_dev = "virtio-gpu")]
register_display_driver!(
    <virtio::VirtIoGpu as VirtIoDevMeta>::Driver,
    <virtio::VirtIoGpu as VirtIoDevMeta>::Device
);

#[cfg(input_dev = "virtio-input")]
register_input_driver!(
    <virtio::VirtIoInput as VirtIoDevMeta>::Driver,
    <virtio::VirtIoInput as VirtIoDevMeta>::Device
);

#[cfg(vsock_dev = "virtio-socket")]
register_vsock_driver!(
    <virtio::VirtIoSocket as VirtIoDevMeta>::Driver,
    <virtio::VirtIoSocket as VirtIoDevMeta>::Device
);

cfg_if::cfg_if! {
    if #[cfg(block_dev = "ramdisk")] {
        pub struct RamDiskDriver;
        register_block_driver!(RamDiskDriver, ax_driver_block::ramdisk::RamDisk);

        impl DriverProbe for RamDiskDriver {
            fn probe_global() -> Option<AxDeviceEnum> {
                // TODO: format RAM disk
                Some(AxDeviceEnum::from_block(
                    ax_driver_block::ramdisk::RamDisk::new(0x100_0000), // 16 MiB
                ))
            }
        }
    }
}

cfg_if::cfg_if! {
    if #[cfg(block_dev = "sdmmc")] {
        pub struct SdMmcDriver;
        register_block_driver!(SdMmcDriver, ax_driver_block::sdmmc::SdMmcDriver);

        impl DriverProbe for SdMmcDriver {
            fn probe_global() -> Option<AxDeviceEnum> {
                let sdmmc = unsafe {
                    ax_driver_block::sdmmc::SdMmcDriver::new(
                        ax_hal::mem::phys_to_virt(ax_config::devices::SDMMC_PADDR.into()).into(),
                    )
                };
                Some(AxDeviceEnum::from_block(sdmmc))
            }
        }
    }
}

cfg_if::cfg_if! {
    if #[cfg(block_dev = "bcm2835-sdhci")]{
        pub struct BcmSdhciDriver;
        register_block_driver!(BcmSdhciDriver, ax_driver_block::bcm2835sdhci::SDHCIDriver);

        impl DriverProbe for BcmSdhciDriver {
            fn probe_global() -> Option<AxDeviceEnum> {
                debug!("mmc probe");
                ax_driver_block::bcm2835sdhci::SDHCIDriver::try_new()
                    .ok()
                    .map(AxDeviceEnum::from_block)
            }
        }
    }
}

cfg_if::cfg_if! {
    if #[cfg(net_dev = "ixgbe")] {
        use crate::ixgbe::IxgbeHalImpl;
        pub struct IxgbeDriver;
        register_net_driver!(IxgbeDriver, ax_driver_net::ixgbe::IxgbeNic<IxgbeHalImpl, 1024, 1>);
        impl DriverProbe for IxgbeDriver {
            #[cfg(bus = "pci")]
            fn probe_pci<C: ax_driver_pci::ConfigurationAccess>(
                root: &mut ax_driver_pci::PciRoot<C>,
                bdf: ax_driver_pci::DeviceFunction,
                dev_info: &ax_driver_pci::DeviceFunctionInfo,
            ) -> Option<crate::AxDeviceEnum> {
                use ax_driver_net::ixgbe::{INTEL_82599, INTEL_VEND, IxgbeNic};
                if dev_info.vendor_id == INTEL_VEND && dev_info.device_id == INTEL_82599 {
                    // Intel 10Gb Network
                    info!("ixgbe PCI device found at {:?}", bdf);

                    // Initialize the device
                    // These can be changed according to the requirments specified in the ixgbe init
                    // function.
                    const QN: u16 = 1;
                    const QS: usize = 1024;
                    let bar_info = root.bar_info(bdf, 0).unwrap();
                    match bar_info {
                        ax_driver_pci::BarInfo::Memory { address, size, .. } => {
                            let ixgbe_nic = IxgbeNic::<IxgbeHalImpl, QS, QN>::init(
                                ax_hal::mem::phys_to_virt((address as usize).into()).into(),
                                size as usize,
                            )
                            .expect("failed to initialize ixgbe device");
                            return Some(AxDeviceEnum::from_net(ixgbe_nic));
                        }
                        ax_driver_pci::BarInfo::IO { .. } => {
                            error!("ixgbe: BAR0 is of I/O type");
                            return None;
                        }
                    }
                }
                None
            }
        }
    }
}

cfg_if::cfg_if! {
    if #[cfg(net_dev = "fxmac")]{
        use ax_alloc::{UsageKind, global_allocator};
        use ax_hal::mem::PAGE_SIZE_4K;

        #[ax_crate_interface::impl_interface]
        impl ax_driver_net::fxmac::KernelFunc for FXmacDriver {
            fn virt_to_phys(addr: usize) -> usize {
                ax_hal::mem::virt_to_phys(addr.into()).into()
            }

            fn phys_to_virt(addr: usize) -> usize {
                ax_hal::mem::phys_to_virt(addr.into()).into()
            }

            fn dma_alloc_coherent(pages: usize) -> (usize, usize) {
                let Ok(vaddr) = global_allocator().alloc_pages(pages, PAGE_SIZE_4K, UsageKind::Dma)
                else {
                    error!("failed to alloc pages");
                    return (0, 0);
                };
                let paddr = ax_hal::mem::virt_to_phys((vaddr).into());
                debug!("alloc pages @ vaddr={:#x}, paddr={:#x}", vaddr, paddr);
                (vaddr, paddr.as_usize())
            }

            fn dma_free_coherent(vaddr: usize, pages: usize) {
                global_allocator().dealloc_pages(vaddr, pages, UsageKind::Dma);
            }

            fn dma_request_irq(_irq: usize, _handler: fn()) {
                warn!("unimplemented dma_request_irq for fxmax");
            }
        }

        register_net_driver!(FXmacDriver, ax_driver_net::fxmac::FXmacNic);

        pub struct FXmacDriver;
        impl DriverProbe for FXmacDriver {
            fn probe_global() -> Option<AxDeviceEnum> {
                info!("fxmac for phytiumpi probe global");
                ax_driver_net::fxmac::FXmacNic::init(0)
                    .ok()
                    .map(AxDeviceEnum::from_net)
            }
        }
    }
}
