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

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

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

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // commented out because it's hard to easily input an IP by command line with docker compose
    // let mut args = std::env::args();
    // args.next();
    // let target = args.next()
    //     .map(|s| s.parse::<Ipv4Addr>().ok())
    //     .flatten()
    //     .ok_or(anyhow!("Please specify a target IP"))?;

    // print!("Please enter the target IP (found with `docker-compose exec examples_revshell_server ip a`, e.x. 172.19.0.2): ");
    // std::io::stdout().flush()?;
    // let stdin = std::io::stdin();
    // let mut target = String::new();
    // stdin.read_line(&mut target)?;
    // let target = target.trim().parse::<Ipv4Addr>()?;
    let target: Ipv4Addr = "172.19.0.2".parse().unwrap();

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

        let routes = socket.get_routes()?;
        let mut routes_inner = routes.iter().collect::<Vec<_>>();
        let neighs = socket.get_neigh()?;
        let links = socket.get_links()?;
        let addrs = socket.get_addrs()?;

        routes_inner.sort_by(|r1, r2| {
            r2.dst().map(|a| a.cidrlen())
                .partial_cmp(&r1.dst().map(|a| a.cidrlen()))
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        println!("-=- Addrs -=-");
        for addr in addrs.iter() {
            println!("addr: {:?}, {}", addr.local(), addr.ifindex());
        }
        println!("-=- Routes -=-");
        for route in routes_inner.iter() {
            println!("route: {:?}", route.dst());

            for hop in route.hop_iter() {
                println!("\thop: {:?}, {}", hop.gateway(), hop.ifindex());
            }
        }
        println!("-=- Neighs -=-");
        for neigh in neighs.iter() {
            println!("neigh: {:?}, {:?}, {}", neigh.dst(), neigh.lladdr(), neigh.ifindex());
        }
        println!("-=- Links -=-");
        for link in links.iter() {
            println!("link {:?}: {:?}, {}", link.name(), link.addr(), link.ifindex());
        }

        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"))?;

        /*let srcip = route::get_srcip_for_dstip(&routes, target)
            .ok_or(anyhow!("Unable to find a route to the IP"))?;

        let (srcip, target_link, dst_mac) = route::get_neigh_for_addr(&routes, &neighs, &links, &srcip.into())
            .ok_or(anyhow!("Unable to find local interface to use"))?;*/

        dbg!(srcmac);
        dbg!(dstmac);
        dbg!(srcip);
        dbg!(target);

        ( 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(())
}