// TODO use more fancy features of virtio-net, TCP/IP checksum offloading would be nice
#![no_std]

extern crate alloc;

use core::mem::{size_of, MaybeUninit};

use alloc::{boxed::Box, sync::Arc};
use bytemuck::{Pod, Zeroable};
use device_api::{
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
    interrupt::{InterruptAffinity, InterruptHandler, IrqVector},
};
use futures_util::task::AtomicWaker;
use libk::{dma::DmaBuffer, task::runtime};
use libk_util::{
    event::BitmapEvent,
    hash_table::DefaultHashTable,
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use ygg_driver_net_core::{
    interface::{NetworkDevice, NetworkInterfaceType},
    RxPacket,
};
use ygg_driver_pci::{
    device::{PciDeviceInfo, PreferredInterruptMode},
    macros::pci_driver,
};
use ygg_driver_virtio_core::{
    queue::{VirtQueue, VqCallbackNotification, VqManualNotification},
    transport::{pci::PciTransport, Transport},
    DeviceStatus,
};
use yggdrasil_abi::{error::Error, net::MacAddress};

pub struct VirtioNet<T: Transport> {
    transport: IrqSafeSpinlock<T>,
    pci_device_info: PciDeviceInfo,
    dma: Arc<dyn DmaAllocator>,

    interface_id: OneTimeInit<u32>,
    mac: IrqSafeRwLock<MacAddress>,

    rx_queue: OneTimeInit<VirtQueue<VqManualNotification>>,
    tx_queue: OneTimeInit<VirtQueue<VqCallbackNotification>>,
    tx_in_flight: IrqSafeSpinlock<DefaultHashTable<u16, DmaBuffer<[u8]>>>,

    softirq: BitmapEvent<AtomicWaker>,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C)]
struct VirtioPacketHeader {
    flags: u8,
    gso_type: u8,
    hdr_len: u16,
    gso_size: u16,
    csum_start: u16,
    csum_offset: u16,
}

impl<T: Transport + 'static> VirtioNet<T> {
    const VQ_RX: u16 = 0;
    const VQ_TX: u16 = 1;

    pub fn new(
        dma: Arc<dyn DmaAllocator>,
        transport: T,
        pci_device_info: PciDeviceInfo,
    ) -> Result<Self, Error> {
        // Read MAC from device config
        let device_cfg = transport
            .device_cfg()
            .expect("virtio-net must have device-specific configuration section");
        let mut mac_bytes = [0; 6];
        mac_bytes.copy_from_slice(&device_cfg[..6]);
        let mac = MacAddress::from(mac_bytes);

        pci_device_info.init_interrupts(PreferredInterruptMode::Msi(true))?;

        Ok(Self {
            transport: IrqSafeSpinlock::new(transport),
            pci_device_info,
            dma,

            interface_id: OneTimeInit::new(),
            mac: IrqSafeRwLock::new(mac),

            rx_queue: OneTimeInit::new(),
            tx_queue: OneTimeInit::new(),
            tx_in_flight: IrqSafeSpinlock::new(DefaultHashTable::new()),
            softirq: BitmapEvent::new(AtomicWaker::new()),
        })
    }

    fn begin_init(&self) -> Result<DeviceStatus, Error> {
        let mut transport = self.transport.lock();
        let mut status = DeviceStatus::RESET_VALUE;

        log::debug!("Reset device");
        transport.write_device_status(status);
        status |= DeviceStatus::ACKNOWLEDGE;
        transport.write_device_status(status);
        status |= DeviceStatus::DRIVER;
        transport.write_device_status(status);

        let _device_features = transport.read_device_features();

        // TODO blah blah blah

        transport.write_driver_features(0);

        status |= DeviceStatus::FEATURES_OK;
        transport.write_device_status(status);

        if !transport
            .read_device_status()
            .contains(DeviceStatus::FEATURES_OK)
        {
            return Err(Error::InvalidOperation);
        }

        Ok(status)
    }

    fn finish_init(&self, status: DeviceStatus) {
        let mut transport = self.transport.lock();

        transport.write_device_status(status | DeviceStatus::DRIVER_OK);
    }

    unsafe fn setup_queues(self: &Arc<Self>) -> Result<(), Error> {
        let (rx_vector, tx_vector) = if let Ok(msis) =
            self.pci_device_info
                .map_interrupt_multiple(0..2, InterruptAffinity::Any, self.clone())
        {
            // Bound a MSI(-x) range, use per-queue vectors
            (Some(msis[0].vector as u16), Some(msis[1].vector as u16))
        } else {
            // TODO support non-MSI-x/non-multivec setups
            self.pci_device_info
                .map_interrupt(InterruptAffinity::Any, self.clone())?;
            (None, None)
        };

        // Setup a callback to remove pending buffers and pass them to the network stack
        let rx_queue = VirtQueue::with_capacity_manual(&*self.dma, Self::VQ_RX, 64)?;

        let p = self.clone();
        // Setup a callback to remove buffers from "in flight owned buffers" list
        let tx_queue = VirtQueue::with_capacity_and_callback(
            &*self.dma,
            Self::VQ_TX,
            64,
            Box::new(move |head, _| {
                p.tx_in_flight.lock().remove(&head);
            }),
        )?;

        let rx_queue = self.rx_queue.init(rx_queue);
        let tx_queue = self.tx_queue.init(tx_queue);

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

        transport.set_queue(Self::VQ_RX, rx_queue, rx_vector);
        transport.set_queue(Self::VQ_TX, tx_queue, tx_vector);

        Ok(())
    }

    async fn softirq(&self) -> Result<(), Error> {
        const RX_SIZE: usize = 4096;
        const RX_IN_FLIGHT: usize = 32;

        let rx_queue = self.rx_queue.get();
        let tx_queue = self.tx_queue.get();
        let nic = *self.interface_id.get();

        let mut rx_in_flight = DefaultHashTable::<u16, DmaBuffer<[MaybeUninit<u8>]>>::new();

        // Setup initial Rx set
        for _ in 0..RX_IN_FLIGHT {
            let mut buffer = DmaBuffer::new_uninit_slice(&*self.dma, RX_SIZE)?;
            let token = rx_queue.try_enqueue(&[], &[buffer.slice_mut(0..RX_SIZE)])?;
            rx_in_flight.insert(token, buffer);
        }

        loop {
            let events = self.softirq.wait().await;

            if events & (1 << Self::VQ_RX) != 0 {
                let refill_rx = rx_queue.handle_notify_manual(|head, _| {
                    if let Some(packet) = rx_in_flight.remove(&head) {
                        let packet = unsafe { DmaBuffer::assume_init_slice(packet) };
                        let packet = RxPacket::new(packet, size_of::<VirtioPacketHeader>(), nic);
                        ygg_driver_net_core::receive_packet(packet).ok();
                    }
                });

                // Refill Rx buffers
                for _ in 0..refill_rx {
                    let mut buffer = DmaBuffer::new_uninit_slice(&*self.dma, RX_SIZE)?;
                    let token = rx_queue.try_enqueue(&[], &[buffer.slice_mut(0..RX_SIZE)])?;
                    rx_in_flight.insert(token, buffer);
                }
            }

            if events & (1 << Self::VQ_TX) != 0 {
                tx_queue.handle_notify();
            }
        }
    }
}

impl<T: Transport + 'static> NetworkDevice for VirtioNet<T> {
    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        DmaBuffer::new_uninit_slice(&*self.dma, len)
    }

    fn transmit_buffer(&self, mut packet: DmaBuffer<[u8]>) -> Result<(), Error> {
        let tx_queue = self.tx_queue.get();
        packet[..size_of::<VirtioPacketHeader>()].fill(0);
        let token = tx_queue.try_enqueue(&[packet.slice(0..packet.len())], &[])?;
        // Add the packet to "in flight" list to make sure it doesn't get dropped and invalidated
        // immediately after returning from this function
        self.tx_in_flight.lock().insert(token, packet);
        self.transport.lock().notify(Self::VQ_TX);
        Ok(())
    }

    fn read_hardware_address(&self) -> MacAddress {
        *self.mac.read()
    }

    fn packet_prefix_size(&self) -> usize {
        size_of::<VirtioPacketHeader>()
    }
}

impl<T: Transport + 'static> InterruptHandler for VirtioNet<T> {
    fn handle_irq(self: Arc<Self>, vector: IrqVector) -> bool {
        match vector {
            IrqVector::Msi(vector) => {
                self.softirq.signal(1 << vector);
                true
            }
            // TODO non-multivec/legacy IRQ setup
            IrqVector::Irq(_) => {
                self.softirq.signal((1 << Self::VQ_RX) | (1 << Self::VQ_TX));
                true
            }
        }
    }
}

impl<T: Transport + 'static> Device for VirtioNet<T> {
    fn display_name(&self) -> &str {
        "VirtIO Network Device"
    }

    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let status = self.begin_init()?;

        self.setup_queues()?;

        self.finish_init(status);

        let iface =
            ygg_driver_net_core::register_interface(NetworkInterfaceType::Ethernet, self.clone());
        self.interface_id.init(iface.id());

        let p = self.clone();
        runtime::spawn(async move { p.softirq().await })?;

        Ok(())
    }

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        Ok(())
    }
}

pci_driver! {
    matches: [device (0x1AF4:0x1000)],
    driver: {
        fn driver_name(&self) -> &str {
            "virtio-net"
        }

        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
            let space = &info.config_space;

            let transport = PciTransport::from_config_space(space).unwrap();
            let device = VirtioNet::new(dma.clone(), transport, info.clone())?;

            let device = Arc::new(device);

            Ok(device)
        }
    }
}
