use alloc::boxed::Box;
use core::{
    format_args,
    pin::Pin,
};
use kernel::{
    bindings,
    c_str,
    driver,
    error::code::*,
    io_mem::IoMem,
    pci, pci::define_pci_id_table,
    prelude::*,
};

use crate::gpu::Gpu;
use crate::gpu::CardType;

struct NovaDevice;

/// BAR Size
pub(crate) const BAR_SIZE: usize = 16777216;

use core::ops::{Add, BitAnd, Div, Not, Sub};

/// Aligns an integer type to a power of two.
pub(crate) fn align<T>(a: T, b: T) -> T
where
    T: Copy
        + Default
        + BitAnd<Output = T>
        + Not<Output = T>
        + Add<Output = T>
        + Sub<Output = T>
        + Div<Output = T>
        + core::cmp::PartialEq,
{
    let def: T = Default::default();
    #[allow(clippy::eq_op)]
    let one: T = !def / !def;

    assert!((b & (b - one)) == def);

    (a + b - one) & !(b - one)
}

impl pci::Driver for NovaDevice {

    define_pci_id_table! {
        (),
        [ (pci::DeviceId::new(bindings::PCI_VENDOR_ID_NVIDIA, bindings::PCI_ANY_ID as u32), None) ]
    }

    fn probe(dev: &mut pci::Device, _id_info: Option<&Self::IdInfo>) -> Result {
        pr_info!("probe()\n");

        dev.enable_device_mem()?;
        dev.set_master();

        let _bars = dev.select_bars(bindings::IORESOURCE_MEM.into());

        let res = dev.take_resource(0).ok_or(ENXIO)?;
        let bar = unsafe { IoMem::<BAR_SIZE>::try_new(res) }?;

        let boot0 = u64::from_le(bar.readq(0));

        if boot0 & 0x1f000000 == 0 {
            return Err(ENODEV);
        }
        pr_info!("nvidia hw rev 0x{:#x}", boot0);
        let chipset = (boot0 & 0x1ff00000) >> 20;
        let _chiprev = boot0 & 0xff;

        let card_type = match chipset & 0x1f0 {
            0x160 => { pr_info!("TU100"); CardType::TU100 },
            0x170 => { pr_info!("GA100"); CardType::GA100 },
            0x190 => { pr_info!("AD100"); CardType::AD100 },
            _ => return Err(ENODEV)
        };

        let mut gpu = Gpu::new(card_type, bar);
        let _ = gpu.init(dev)?;

        Ok(())
    }

    fn remove(_data: &Self::Data) {
        pr_info!("remove()\n");
    }
}

pub(crate) struct NovaModule {
    _registration: Pin<Box<driver::Registration<pci::Adapter<NovaDevice>>>>,
}

impl kernel::Module for NovaModule {
    fn init(_name: &'static CStr, module: &'static ThisModule) -> Result<Self> {
        pr_info!("Module loaded!\n");

        let registration = driver::Registration::new_pinned(c_str!("nova"), module)?;
        pr_info!("Registerd PCI driver.\n");

        Ok(Self {
            _registration: registration,
        })
    }
}