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

use alloc::{sync::Arc, vec::Vec};
use bytemuck::{Pod, Zeroable};
use libk::dma::DmaBuffer;
use libk_util::sync::spin_rwlock::{
    IrqSafeRwLock, IrqSafeRwLockReadGuard, IrqSafeRwLockWriteGuard,
};
use yggdrasil_abi::{
    error::Error,
    net::{
        protocols::{EtherType, EthernetFrame, InetChecksum, IpProtocol, Ipv4Frame},
        types::NetValueImpl,
        MacAddress, SubnetAddr,
    },
};

use crate::{interface::NetworkInterface, l4, TxPacketBuilder};

pub mod arp;
pub mod ip;

pub struct L3Packet {
    pub interface_id: u32,

    pub protocol: IpProtocol,

    pub source_address: IpAddr,
    pub destination_address: IpAddr,

    pub source_port: Option<u16>,
    pub destination_port: Option<u16>,

    pub l2_offset: usize,
    pub l3_offset: usize,
    pub l4_offset: usize,
    pub data_length: usize,

    pub data: Arc<DmaBuffer<[u8]>>,
}

pub trait IpFrame: Pod {
    fn destination_ip(&self) -> IpAddr;
    fn source_ip(&self) -> IpAddr;
    fn data_length(&self) -> usize;
}

// TODO use range map for this?
pub struct Route {
    pub subnet: SubnetAddr,
    pub interface: u32,
    pub gateway: Option<IpAddr>,
    pub priority: u32,
}

pub struct L4ResolvedPacket<'a, 'i> {
    pub interface: &'i NetworkInterface,
    pub source_ip: IpAddr,
    pub gateway_ip: IpAddr,
    pub destination_ip: IpAddr,
    pub l4_frame: &'a [u8],
    pub l4_options: &'a [u8],
    pub l4_data: &'a [u8],
    pub protocol: IpProtocol,
    pub ttl: u8,
}

pub struct L4UnresolvedPacket<'a> {
    pub destination_ip: IpAddr,
    pub l4_frame: &'a [u8],
    pub l4_options: &'a [u8],
    pub l4_data: &'a [u8],
    pub protocol: IpProtocol,
    pub ttl: u8,
}

static ROUTES: IrqSafeRwLock<Vec<Route>> = IrqSafeRwLock::new(Vec::new());

impl L3Packet {
    pub fn l3_data(&self) -> &[u8] {
        &self.data[self.l4_offset..]
    }
}

impl Route {
    pub fn list_mut() -> IrqSafeRwLockWriteGuard<'static, Vec<Route>> {
        ROUTES.write()
    }

    pub fn list_ref() -> IrqSafeRwLockReadGuard<'static, Vec<Route>> {
        ROUTES.read()
    }

    pub fn lookup(address: IpAddr) -> Option<(u32, Option<IpAddr>, IpAddr)> {
        // TODO sort routes based on their "specificity"?
        // Check for local route
        for (_, interface) in NetworkInterface::list_ref().iter() {
            if interface
                .address
                .read()
                .map(|addr| addr == address)
                .unwrap_or(false)
            {
                // This is the address of loopback, return it
                return Some((
                    NetworkInterface::loopback().id,
                    Some(IpAddr::V4(Ipv4Addr::LOCALHOST)),
                    IpAddr::V4(Ipv4Addr::LOCALHOST),
                ));
            }
        }

        let routes = ROUTES.read();
        let mut best: Option<(&Route, IpAddr)> = None;
        for route in routes.iter() {
            if route.subnet.contains(&address) || route.subnet.base() == address {
                if let Some((prev_best, _)) = best.as_ref() {
                    if route.priority < prev_best.priority {
                        best = Some((route, address));
                    }
                } else {
                    best = Some((route, address));
                }
            }
        }
        let (best, address) = best?;
        Some((best.interface, best.gateway, address))
    }

    pub fn insert(route: Self) -> Result<(), Error> {
        // TODO check for conflicts
        log::debug!("Add route: {route}");
        ROUTES.write().push(route);
        Ok(())
    }
}

impl fmt::Display for Route {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} ", self.subnet)?;
        if let Some(gw) = self.gateway {
            write!(f, " via {gw}")?;
        }
        Ok(())
    }
}

pub async fn handle_accepted(l3_packet: L3Packet) -> Result<(), Error> {
    match l3_packet.protocol {
        IpProtocol::UDP => l4::udp::handle(l3_packet),
        IpProtocol::ICMP => l4::icmp::handle(l3_packet).await,
        IpProtocol::TCP => l4::tcp::handle(l3_packet).await,
        _ => Ok(()),
    }
}

impl L4ResolvedPacket<'_, '_> {
    pub async fn lookup_gateway_mac(&self) -> Result<MacAddress, Error> {
        arp::lookup(self.interface.id, self.gateway_ip, true).await
    }

    pub fn make_l3_frame(&self) -> Result<Ipv4Frame, Error> {
        // TODO what if source_ip/gateway_ip/destination_ip are a mix of IPv4/IPv6?
        let source_ip = match self.source_ip {
            IpAddr::V4(v4) => v4,
            IpAddr::V6(_) => {
                return Err(Error::NotImplemented);
            }
        };
        let destination_ip = match self.destination_ip {
            IpAddr::V4(v4) => v4,
            IpAddr::V6(_) => {
                return Err(Error::NotImplemented);
            }
        };

        let total_length = (self.total_l4_len() + size_of::<Ipv4Frame>())
            .try_into()
            .map_err(|_| Error::InvalidArgument)?;

        let mut l3_frame = Ipv4Frame {
            source_address: u32::to_network_order(source_ip.into()),
            destination_address: u32::to_network_order(destination_ip.into()),
            protocol: self.protocol,
            version_length: 0x45,
            total_length: u16::to_network_order(total_length),
            flags_frag: u16::to_network_order(0x4000),
            id: u16::to_network_order(0),
            ttl: self.ttl,
            ..Ipv4Frame::zeroed()
        };

        let l3_frame_bytes = bytemuck::bytes_of(&l3_frame);
        let mut ip_checksum = InetChecksum::new();
        ip_checksum.add_bytes(l3_frame_bytes, true);
        let ip_checksum = ip_checksum.finish();

        l3_frame.header_checksum = u16::to_network_order(ip_checksum);

        Ok(l3_frame)
    }

    pub fn total_l4_len(&self) -> usize {
        self.l4_frame.len() + self.l4_options.len() + self.l4_data.len()
    }
}

pub(crate) fn resolve_l3_route(
    destination_ip: IpAddr,
) -> Result<(Arc<NetworkInterface>, IpAddr, IpAddr, IpAddr), Error> {
    // Find the route itself
    let (interface_id, gateway, destination_ip) =
        Route::lookup(destination_ip).ok_or(Error::NetworkUnreachable)?;
    let interface = NetworkInterface::get(interface_id)?;
    // Route exists, but has no gateway (TODO: assign subnets to interfaces)
    let gateway = gateway.ok_or(Error::NetworkUnreachable)?;
    // Route exists, but network has no address assigned (TODO: how?)
    let source_address = interface.address.read().ok_or(Error::NetworkUnreachable)?;

    Ok((interface, source_address, gateway, destination_ip))
}

pub async fn send_l4_ip_resolved(packet: &L4ResolvedPacket<'_, '_>) -> Result<(), Error> {
    let gateway_mac = packet.lookup_gateway_mac().await?;

    let l3_frame = packet.make_l3_frame()?;

    let mut builder = TxPacketBuilder::new(
        packet.interface,
        size_of::<EthernetFrame>() + size_of::<Ipv4Frame>() + packet.total_l4_len(),
    )?;
    builder.push(&EthernetFrame {
        source_mac: packet.interface.mac,
        destination_mac: gateway_mac,
        ethertype: EtherType::IPV4.to_network_order(),
    })?;
    builder.push(&l3_frame)?;
    builder.push_bytes(packet.l4_frame)?;
    builder.push_bytes(packet.l4_options)?;
    builder.push_bytes(packet.l4_data)?;
    builder.transmit(None)
}

pub async fn send_l4_ip(packet: &L4UnresolvedPacket<'_>) -> Result<(), Error> {
    let (interface, source_ip, gateway_ip, destination_ip) =
        resolve_l3_route(packet.destination_ip)?;

    send_l4_ip_resolved(&L4ResolvedPacket {
        interface: &interface,
        source_ip,
        gateway_ip,
        destination_ip,
        l4_frame: packet.l4_frame,
        l4_options: packet.l4_options,
        l4_data: packet.l4_data,
        protocol: packet.protocol,
        ttl: packet.ttl,
    })
    .await
}
