#![feature(map_try_insert, let_chains)]
#![allow(clippy::type_complexity, clippy::new_without_default)]
#![no_std]

extern crate alloc;

use core::mem::size_of;

use alloc::sync::Arc;
use bytemuck::Pod;
use ethernet::L2Packet;
use interface::NetworkInterface;
use l3::L3Packet;
use libk::{dma::DmaBuffer, task::runtime};
use libk_util::queue::UnboundedMpmcQueue;
use socket::RawSocket;
use yggdrasil_abi::{error::Error, net::protocols::EthernetFrame};

pub mod ethernet;
pub mod l3;
pub mod l4;

pub mod socket;

pub mod interface;
pub mod util;

pub mod ephy;

pub use interface::register_interface;

pub struct RxPacket {
    buffer: DmaBuffer<[u8]>,
    offset: usize,
    iface: u32,
}

pub struct TxPacketBuilder<'a> {
    nic: &'a NetworkInterface,
    data: DmaBuffer<[u8]>,
    pos: usize,
    #[allow(unused)]
    l2_offset: usize,
}

impl<'a> TxPacketBuilder<'a> {
    pub fn new(nic: &'a NetworkInterface, l2_size: usize) -> Result<Self, Error> {
        let l2_offset = nic.device.packet_prefix_size();
        let data = nic.device.allocate_transmit_buffer(l2_offset + l2_size)?;
        let data = unsafe { DmaBuffer::assume_init_slice(data) };
        Ok(Self {
            nic,
            data,
            pos: l2_offset,
            l2_offset,
        })
    }

    #[inline]
    pub fn push<T: Pod>(&mut self, value: &T) -> Result<(), Error> {
        self.push_bytes(bytemuck::bytes_of(value))
    }

    pub fn push_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> {
        if self.pos + bytes.len() > self.data.len() {
            return Err(Error::OutOfMemory);
        }
        self.data[self.pos..self.pos + bytes.len()].copy_from_slice(bytes);
        self.pos += bytes.len();
        Ok(())
    }

    fn finish(self) -> (DmaBuffer<[u8]>, usize) {
        (self.data, self.pos)
    }

    pub fn transmit(self, mute_raw: Option<u32>) -> Result<(), Error> {
        let nic = self.nic;
        let (data, len) = self.finish();
        // Also transmit to raw
        RawSocket::handle_tx(nic.id, &data, len, mute_raw);
        nic.device.transmit_buffer(data)
    }
}

impl RxPacket {
    #[inline]
    pub fn new(buffer: DmaBuffer<[u8]>, offset: usize, iface: u32) -> Self {
        Self {
            buffer,
            offset,
            iface,
        }
    }
}

static PACKET_QUEUE: UnboundedMpmcQueue<RxPacket> = UnboundedMpmcQueue::new();
static ACCEPT_QUEUE: UnboundedMpmcQueue<L3Packet> = UnboundedMpmcQueue::new();

#[inline]
pub fn receive_packet(packet: RxPacket) -> Result<(), Error> {
    PACKET_QUEUE.push_back(packet);
    Ok(())
}

pub fn start_network_tasks() -> Result<(), Error> {
    runtime::spawn(l2_packet_handler_worker())?;
    for _ in 0..4 {
        runtime::spawn(l3_accept_worker())?;
    }

    Ok(())
}

async fn l2_packet_handler_worker() {
    loop {
        let packet = PACKET_QUEUE.pop_front().await;

        let eth_frame: &EthernetFrame = bytemuck::from_bytes(
            &packet.buffer[packet.offset..packet.offset + size_of::<EthernetFrame>()],
        );

        let l2_packet = L2Packet {
            interface_id: packet.iface,

            source_address: eth_frame.source_mac,
            destination_address: eth_frame.destination_mac,

            l2_offset: packet.offset,
            l3_offset: packet.offset + size_of::<EthernetFrame>(),

            data: Arc::new(packet.buffer),
        };

        ethernet::handle(l2_packet);
    }
}

async fn l3_accept_worker() {
    loop {
        let l3_packet = ACCEPT_QUEUE.pop_front().await;

        // log::debug!(
        //     "INPUT {} {}:{:?} -> {}:{:?}: ACCEPT",
        //     l3_packet.protocol,
        //     l3_packet.source_address,
        //     l3_packet.source_port,
        //     l3_packet.destination_address,
        //     l3_packet.destination_port
        // );

        if let Err(error) = l3::handle_accepted(l3_packet).await {
            log::error!("L3 handle error: {:?}", error);
        }
    }
}
