use core::{
    mem::size_of,
    net::{IpAddr, Ipv4Addr, SocketAddr},
};

use libk::error::Error;
use yggdrasil_abi::{
    abi_serde::wire,
    net::{
        protocols::{IpProtocol, Ipv4Frame, TcpFrame, UdpFrame},
        types::NetValueImpl,
    },
};

use crate::{interface::NetworkInterface, L2Packet, L3Packet, ACCEPT_QUEUE};

use super::IpFrame;

impl IpFrame for Ipv4Frame {
    fn destination_ip(&self) -> IpAddr {
        IpAddr::V4(Ipv4Addr::from(u32::from_network_order(
            self.destination_address,
        )))
    }

    fn source_ip(&self) -> IpAddr {
        IpAddr::V4(Ipv4Addr::from(u32::from_network_order(self.source_address)))
    }

    fn data_length(&self) -> usize {
        self.total_length().saturating_sub(self.header_length())
    }
}

pub fn load_address(bytes: &[u8]) -> Result<SocketAddr, Error> {
    Ok(wire::from_slice(bytes)?)
}
pub fn store_address(source: &mut Option<&mut [u8]>, address: &SocketAddr) -> Result<(), Error> {
    if let Some(buffer) = source.as_mut() {
        wire::to_slice(address, buffer)?;
    }
    Ok(())
}

pub fn handle_v4_packet(packet: L2Packet) {
    let Ok(interface) = NetworkInterface::get(packet.interface_id) else {
        log::debug!("Invalid interface ID in L2 packet");
        return;
    };

    let l2_data = packet.l2_data();
    let l3_frame: &Ipv4Frame = bytemuck::from_bytes(&l2_data[..size_of::<Ipv4Frame>()]);
    let header_length = l3_frame.header_length();
    let l3_data = &l2_data[size_of::<Ipv4Frame>()..];

    let is_input = interface
        .address
        .read()
        .map(|address| address == l3_frame.destination_ip())
        .unwrap_or(false);

    if is_input {
        // Extract ports from L4 proto
        let (source_port, destination_port) = match l3_frame.protocol {
            IpProtocol::UDP => {
                // TODO check size
                let l4_frame: &UdpFrame = bytemuck::from_bytes(&l3_data[..size_of::<UdpFrame>()]);
                (
                    Some(u16::from_network_order(l4_frame.source_port)),
                    Some(u16::from_network_order(l4_frame.destination_port)),
                )
            }
            IpProtocol::TCP => {
                // TODO check size
                let l4_frame: &TcpFrame = bytemuck::from_bytes(&l3_data[..size_of::<TcpFrame>()]);
                (
                    Some(u16::from_network_order(l4_frame.source_port)),
                    Some(u16::from_network_order(l4_frame.destination_port)),
                )
            }
            IpProtocol::ICMP => (None, None),
            _ => (None, None),
        };

        let l3_packet = L3Packet {
            interface_id: packet.interface_id,

            protocol: l3_frame.protocol,

            source_address: l3_frame.source_ip(),
            destination_address: l3_frame.destination_ip(),

            source_port,
            destination_port,

            l2_offset: packet.l2_offset,
            l3_offset: packet.l3_offset,
            l4_offset: packet.l3_offset + header_length,
            data_length: l3_frame.data_length(),

            data: packet.data,
        };

        ACCEPT_QUEUE.push_back(l3_packet);
    } else {
        // TODO forwarding
        log::debug!(
            "Dropped forwarded IPv4: {} -> {}",
            l3_frame.source_ip(),
            l3_frame.destination_ip()
        );
    }
}
