#![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 futures_util::task::AtomicWaker;
use libk::{dma::DmaBuffer, error::Error, task::runtime};
use libk_util::{OneTimeInit, event::BitmapEvent, sync::IrqSafeSpinlock};
use regs::Regs;
use tock_registers::{LocalRegisterCopy, fields::FieldValue};
use ygg_driver_net_core::{
    RxPacket,
    interface::{NetworkDevice, NetworkInterfaceType},
    register_interface,
    util::GenericQueue,
};
use ygg_driver_pci::{
    PciBaseAddress, PciConfigurationSpace,
    device::{PciDeviceInfo, PreferredInterruptMode},
    macros::pci_driver,
};
use yggdrasil_abi::net::{MacAddress, link::LinkState};

use crate::{
    regs::Revision,
    ring::{RxDescriptor, TxDescriptor},
};

extern crate alloc;

mod regs;
mod ring;

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

    mac: OneTimeInit<MacAddress>,
    queue: OneTimeInit<GenericQueue<TxDescriptor, RxDescriptor>>,
    softirq: BitmapEvent<AtomicWaker>,
    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),

            queue: OneTimeInit::new(),
            softirq: BitmapEvent::new(AtomicWaker::new()),
            nic: OneTimeInit::new(),
        }
    }

    async fn softirq(self: Arc<Self>) {
        let queue = self.queue.get();
        let rx_capacity = queue.rx_ring.capacity();
        let nic = *self.nic.get();

        loop {
            let event = self.softirq.wait().await as u32;
            let event = LocalRegisterCopy::<_, regs::ICR::Register>::new(event);

            if event.matches_all(regs::ICR::LSC::SET) {
                let status = self.regs.lock().read_link();
                log::info!("igbe: link is {status}");
                self.interrupt_handled(regs::ICR::LSC::SET);
            }
            if event.matches_any(&[regs::ICR::TXQE::SET, regs::ICR::TXDW::SET]) {
                let head = self.regs.lock().tx_queue_head();
                queue.drop_tx_until(head as usize);
                // queue.consume_tx(Some(head as usize)).ok();
                self.interrupt_handled(regs::ICR::TXQE::SET + regs::ICR::TXDW::SET);
            }
            if event.matches_any(&[regs::ICR::RXDMT0::SET, regs::ICR::RXT0::SET]) {
                {
                    let mut regs = self.regs.lock();
                    let head = regs.rx_queue_head() as usize;
                    let tail = queue
                        .consume_rx(&*self.dma, Some(head), |packet, _| {
                            let packet = RxPacket::new(packet, 0, nic);
                            ygg_driver_net_core::receive_packet(packet).ok();
                        })
                        .expect("Rx ring handle error");
                    let tail = (tail + rx_capacity - 1) & (rx_capacity - 1);
                    regs.set_rx_queue_tail(tail as u16);
                }
                self.interrupt_handled(regs::ICR::RXDMT0::SET + regs::ICR::RXT0::SET);
            }
        }
    }

    fn interrupt_handled(&self, icr: FieldValue<u32, regs::ICR::Register>) {
        let mut regs = self.regs.lock();
        regs.clear_interrupts(icr.value);
        regs.enable_interrupts(icr.value);
    }
}

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 queue = self.queue.init(GenericQueue::with_capacity(
            &*self.dma,
            128,
            128,
            2048 + 16,
        )?);

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

        regs.disable_interrupts(0xFFFFFFFF);
        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(0xFFFFFFFF);
        regs.set_link_up(self.chip)?;

        // Initialize Rx
        regs.initialize_receiver(&queue.rx_ring);
        regs.initialize_transmitter(&queue.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);
        }

        let nic = register_interface(NetworkInterfaceType::Ethernet, self.clone());
        self.nic.init(nic.id());

        regs.enable_interrupts(
            (regs::IMS::TXDW::SET
                + regs::IMS::TXQE::SET
                + regs::IMS::LSC::SET
                + regs::IMS::RXT0::SET
                + regs::IMS::RXDMT0::SET)
                .value,
        );

        runtime::spawn(self.clone().softirq())?;

        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().get();
        if cause == 0 {
            return false;
        }
        regs.disable_interrupts(cause);
        self.softirq.signal(cause as u64);
        true
    }
}

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 queue = self.queue.get();
        let tail = queue.try_push_xmit(buffer)?;
        let mut regs = self.regs.lock();
        regs.set_tx_queue_tail(tail as _);
        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"
        }
    }
}
