#![no_std]

use core::{mem::MaybeUninit, time::Duration};

use alloc::sync::Arc;
use device_api::{
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
    interrupt::{InterruptAffinity, InterruptHandler, IrqVector},
};
use libk::{dma::DmaBuffer, error::Error};
use libk_util::{
    sync::{IrqSafeSpinlock, Spinlock},
    OneTimeInit,
};
use regs::{Regs, ICR};
use ring::{RxRing, TxRing};
use ygg_driver_net_core::{
    interface::{NetworkDevice, NetworkInterfaceType},
    register_interface, RxPacket,
};
use ygg_driver_pci::{
    device::{PciDeviceInfo, PreferredInterruptMode},
    macros::pci_driver,
    PciBaseAddress, PciConfigurationSpace,
};
use yggdrasil_abi::net::{link::LinkState, MacAddress};

use crate::regs::Revision;

extern crate alloc;

mod regs;
mod ring;

struct Igbe {
    chip: Revision,
    regs: IrqSafeSpinlock<Regs>,
    dma: Arc<dyn DmaAllocator>,
    pci: PciDeviceInfo,

    mac: OneTimeInit<MacAddress>,
    rx_ring: OneTimeInit<Spinlock<RxRing>>,
    tx_ring: OneTimeInit<IrqSafeSpinlock<TxRing>>,
    nic: OneTimeInit<u32>,
}

impl Igbe {
    pub fn new(dma: Arc<dyn DmaAllocator>, regs: Regs, chip: Revision, pci: PciDeviceInfo) -> Self {
        Self {
            chip,
            dma,
            pci,
            mac: OneTimeInit::new(),
            regs: IrqSafeSpinlock::new(regs),

            rx_ring: OneTimeInit::new(),
            tx_ring: OneTimeInit::new(),
            nic: OneTimeInit::new(),
        }
    }
}

impl Device for Igbe {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let msi_info = self
            .pci
            .map_interrupt(InterruptAffinity::Any, self.clone())?;

        let rx_ring = RxRing::with_capacity(&*self.dma, 128, 2048 + 16)?;
        let tx_ring = TxRing::with_capacity(&*self.dma, 128)?;

        let mut regs = self.regs.lock();

        regs.disable_interrupts();
        let mac = regs.read_mac()?;
        self.mac.init(mac);
        regs.reset(Duration::from_millis(200))?;
        // Intel 8257x manuals say an additional interrupt disable is needed after a global reset
        regs.disable_interrupts();
        regs.set_link_up(self.chip)?;

        // Initialize Rx
        regs.initialize_receiver(&rx_ring);
        regs.initialize_transmitter(&tx_ring);
        // If MSI(-x) was initialized, notify the NIC about it
        if let Some(msi_info) = msi_info {
            regs.initialize_ivar(msi_info.vector);
        }

        self.rx_ring.init(Spinlock::new(rx_ring));
        self.tx_ring.init(IrqSafeSpinlock::new(tx_ring));
        let nic = register_interface(NetworkInterfaceType::Ethernet, self.clone());
        self.nic.init(nic.id());

        regs.enable_interrupts();

        Ok(())
    }

    fn display_name(&self) -> &str {
        "Intel Gigabit Ethernet"
    }
}

impl InterruptHandler for Igbe {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let mut regs = self.regs.lock();
        let cause = regs.interrupt_cause();
        if cause.get() == 0 {
            return false;
        }
        regs.clear_interrupts(cause.get());
        let mut any = false;

        if cause.matches_all(ICR::LSC::SET) {
            let status = regs.read_link();
            log::info!("igbe: link is {status}");
            any = true;
        }

        if cause.matches_all(ICR::RXT0::SET) {
            let mut rx = self.rx_ring.get().lock();
            let nic = *self.nic.get();
            let head = regs.rx_queue_head();
            let tail = rx.handle_rx(&*self.dma, head, |packet, _| {
                let packet = RxPacket::new(packet, 0, nic);
                ygg_driver_net_core::receive_packet(packet).ok();
            });
            regs.set_rx_queue_tail(tail);
            any = true;
        }

        if cause.matches_any(&[ICR::TXQE::SET, ICR::TXDW::SET]) {
            let mut tx = self.tx_ring.get().lock();
            let head = regs.tx_queue_head();
            tx.handle_tx(head);
            any = true;
        }

        if !any {
            log::info!("igbe: unhandled irq {:#x}", cause.get());
        }

        any
    }
}

impl NetworkDevice for Igbe {
    fn read_hardware_address(&self) -> MacAddress {
        *self.mac.get()
    }

    fn packet_prefix_size(&self) -> usize {
        0
    }

    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        DmaBuffer::new_uninit_slice(&*self.dma, len + 4)
    }

    fn transmit_buffer(&self, buffer: DmaBuffer<[u8]>) -> Result<(), Error> {
        let mut tx = self.tx_ring.get().lock();

        let Ok(head) = tx.tx_now(buffer) else {
            return Err(Error::WouldBlock);
        };

        let mut regs = self.regs.lock();
        regs.set_tx_queue_tail(head);

        Ok(())
    }

    fn link_state(&self) -> LinkState {
        LinkState::Ethernet(self.regs.lock().read_link())
    }
}

pci_driver! {
    matches: [
        device (0x8086:0x100E), // 82540EM (E1000)
        device (0x8086:0x100C), // 82544GC (E1000)
        device (0x8086:0x100F), // 82545EM (E1000)
        device (0x8086:0x10D3), // 82574L (E1000E) [[BROKEN]]
        device (0x8086:0x10C9), // 82576 GbE
        device (0x8086:0x1502), // 82579LM GbE (Lewisville)
    ],
    driver: {
        fn probe(
            &self,
            info: &PciDeviceInfo,
            dma: &Arc<dyn DmaAllocator>,
        ) -> Result<Arc<dyn Device>, Error> {
            let base = info
                .config_space
                .bar(0).ok_or(Error::InvalidArgument)?;

            let use_mmio = match base {
                PciBaseAddress::Memory32(_) | PciBaseAddress::Memory64(_) => {
                    true
                }
                PciBaseAddress::Io(_) => {
                    false
                }
            };

            let chip = match info.device_id {
                0x100E | 0x100C | 0x100F => Revision::I8254x,
                0x10D3 => Revision::I82574L,
                0x10C9 => Revision::I82576,
                0x1502 => Revision::I82579LM,
                id => {
                    log::error!("Invalid igbe chip variant: {id:#04x}");
                    return Err(Error::InvalidOperation)
                },
            };

            info.init_interrupts(PreferredInterruptMode::Msi(true))?;
            info.set_command(true, use_mmio, !use_mmio, true);

            let regs = unsafe { Regs::map(base) }?;
            let device = Igbe::new(dma.clone(), regs, chip, info.clone());
            Ok(Arc::new(device))
        }

        fn driver_name(&self) -> &str {
            "igbe"
        }
    }
}
