use std::{collections::HashMap, net::Ipv4Addr};

use anyhow::anyhow;
use tokio::time::{interval, Duration};
use tokio_stream::StreamExt;

use nl_sys::{netlink, route};
use pcap_sys::packets::*;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let target = std::env::args().skip(1).next().ok_or(anyhow!("could not get target IP"))?.parse::<Ipv4Addr>()?;

    let (ifname, src_mac, dst_mac, srcip) = {
        let socket = netlink::Socket::new()?;

        let routes = socket.get_routes()?;
        let neighs = socket.get_neigh()?;
        let links = socket.get_links()?;
        let addrs = socket.get_addrs()?;

        let (ifname, srcip, srcmac, dstmac) =
            route::get_macs_and_src_for_ip(&addrs, &routes, &neighs, &links, target)
                .ok_or(anyhow!("unable to find a route to the IP"))?;

        (ifname, srcmac, dstmac, srcip)
    };

    let mut interface = pcap_sys::new_aggregate_interface(false)?;

    interface.set_buffer_size(8192)?;
    interface.set_non_blocking(true)?;
    interface.set_promisc(false)?;
    interface.set_timeout(10)?;

    let mut interface = interface.activate()?;

    interface.set_filter("inbound and udp port 54248", true, None)?;

    interface.prune(|_, interface| interface.datalink() != pcap_sys::consts::DLT_EN10MB);

    enum EventType {
        Packet((String, Result<EthernetPacket, pcap_sys::error::Error>)),
        Update,
    }

    let mut packets = interface.stream()?;
    let mut update_interval = interval(Duration::from_millis(500));
    let mut current_packet_id = 0;
    let mut sent_updates: HashMap<i32, bool> = HashMap::new();

    while let Some(evt) = tokio::select! {
        v = packets.next() => v.map(EventType::Packet),
        _ = update_interval.tick() => Some(EventType::Update)
    } {
        match evt {
            EventType::Packet((_, Ok(pkt))) => {
                let eth_pkt = pkt.pkt();
                let Ok(Layer3Pkt::IPv4Pkt(ip_pkt)) = eth_pkt.get_layer3_pkt() else { continue; };
                let Ok(Layer4Pkt::UDP(udp_pkt)) = ip_pkt.get_layer4_packet() else { continue; };

                let data = udp_pkt.get_data();

                let Ok(resp_id) = TryInto::<[u8;4]>::try_into(&data[..4]) else { continue; };
                let resp_id = i32::from_be_bytes(resp_id);

                if sent_updates.contains_key(&resp_id) {
                    sent_updates.remove(&resp_id);
                    println!("Packet {resp_id} has received a response");
                }
            }
            EventType::Update => {
                let unacknowledged_packets = sent_updates.keys().collect::<Vec<_>>();
                println!("Currently unacknowledged packets: {unacknowledged_packets:?}");
                current_packet_id += 1;
                sent_updates.insert(current_packet_id, false);

                let udp_packet =
                    UDPPacket::construct(54248, 54248, current_packet_id.to_be_bytes().to_vec());
                let ip_packet =
                    IPv4Packet::construct(srcip, target, &Layer4Packet::UDP(udp_packet));
                let eth_packet =
                    EthernetPacket::construct(src_mac, dst_mac, &Layer3Packet::IPv4(ip_packet));

                packets.sendpacket(&ifname, eth_packet.pkt())?;
            }
            _ => {}
        }
    }

    Ok(())
}