#![feature(yggdrasil_os, rustc_private)]

use std::{
    mem::size_of,
    net::{IpAddr, Ipv4Addr, ToSocketAddrs},
    os::{
        fd::AsRawFd,
        yggdrasil::io::{net::raw_socket::RawSocket, poll::PollChannel, timer::TimerFd},
    },
    process::ExitCode,
    str::FromStr,
    sync::atomic::{AtomicBool, Ordering},
    time::Duration,
};

use bytemuck::Zeroable;
use clap::Parser;
use netutils::{netconfig::NetConfig, Error};
use runtime::abi::net::MacAddress;
use yggdrasil_abi::net::{
    protocols::{EtherType, EthernetFrame, IcmpV4Frame, InetChecksum, IpProtocol, Ipv4Frame},
    types::NetValueImpl,
};

#[derive(Parser)]
struct Args {
    #[clap(
        help = "Time (ms) between a reply is received and the next request is sent",
        short,
        long,
        default_value_t = 1000,
        value_parser = valid_interval
    )]
    inteval: u32,
    #[clap(
        help = "Time (ms) after which the request is considered unanswered",
        short,
        long,
        default_value_t = 500,
        value_parser = valid_timeout,
    )]
    timeout: u32,
    #[clap(
        help = "Number of requests to perform",
        short,
        long,
        default_value_t = 10
    )]
    count: usize,
    #[clap(
        help = "Amount of bytes to include as data",
        short,
        long,
        default_value_t = 64,
        value_parser = valid_data_size
    )]
    data_size: usize,

    #[clap(help = "Address to ping")]
    address: String,
}

fn valid_interval(s: &str) -> Result<u32, String> {
    clap_num::number_range(s, 100, 10000)
}

fn valid_timeout(s: &str) -> Result<u32, String> {
    clap_num::number_range(s, 100, 5000)
}

fn valid_data_size(s: &str) -> Result<usize, String> {
    clap_num::number_range(s, 4, 128)
}

struct PingRouting {
    interface_id: u32,
    source_ip: IpAddr,
    destination_ip: IpAddr,
    source_mac: MacAddress,
    gateway_mac: MacAddress,
}

struct PingStats {
    packets_sent: usize,
    packets_received: usize,
}

fn resolve_routing(address: IpAddr) -> Result<PingRouting, Error> {
    let mut nc = NetConfig::open()?;
    let routing = nc.query_route(address)?;
    let Some(source) = routing.source else {
        todo!();
    };
    let Some(gateway) = routing.gateway else {
        todo!();
    };

    let gateway_mac = nc.query_arp(routing.interface_id, gateway, false)?;

    Ok(PingRouting {
        interface_id: routing.interface_id,
        source_ip: source,
        destination_ip: routing.destination,
        source_mac: routing.source_mac,
        gateway_mac,
    })
}

fn validate_ping_reply(
    packet: &[u8],
    local: Ipv4Addr,
    remote: Ipv4Addr,
    expect_l4_data: &[u8],
    expect_id: u16,
    expect_seq: u16,
) -> bool {
    if packet.len() < size_of::<EthernetFrame>() + size_of::<Ipv4Frame>() {
        return false;
    }

    let l3_offset = size_of::<EthernetFrame>();

    let l2_frame: &EthernetFrame = bytemuck::from_bytes(&packet[..l3_offset]);

    if EtherType::from_network_order(l2_frame.ethertype) != EtherType::IPV4 {
        return false;
    }
    let l3_frame: &Ipv4Frame =
        bytemuck::from_bytes(&packet[l3_offset..l3_offset + size_of::<Ipv4Frame>()]);
    if l3_frame.protocol != IpProtocol::ICMP
        || u32::from_network_order(l3_frame.source_address) != u32::from(remote)
        || u32::from_network_order(l3_frame.destination_address) != u32::from(local)
    {
        return false;
    }
    let mut ip_checksum = InetChecksum::new();
    ip_checksum.add_value(l3_frame, true);
    let ip_checksum = ip_checksum.finish();

    if ip_checksum != 0 {
        eprintln!("IP checksum mismatch: {:#06x}", ip_checksum);
        return false;
    }

    let l4_offset = l3_offset + l3_frame.header_length();
    let l4_size = l3_frame
        .total_length()
        .saturating_sub(l3_frame.header_length());
    if packet.len() < l4_offset + size_of::<IcmpV4Frame>() + expect_l4_data.len() {
        return false;
    }
    let l4_frame: &IcmpV4Frame =
        bytemuck::from_bytes(&packet[l4_offset..l4_offset + size_of::<IcmpV4Frame>()]);
    let l4_data = &packet[l4_offset + size_of::<IcmpV4Frame>()..l4_offset + l4_size];

    if l4_frame.ty != 0 || l4_frame.code != 0 {
        return false;
    }

    let rest = u32::from_network_order(l4_frame.rest);
    let reply_id = (rest >> 16) as u16;
    let reply_seq = rest as u16;

    if reply_id != expect_id || reply_seq != expect_seq {
        eprintln!(
            "ICMP seq/id mismatch: sent {}/{}, got {}/{}",
            expect_id, expect_seq, reply_id, reply_seq
        );
        return false;
    }

    let mut icmp_checksum = InetChecksum::new();
    icmp_checksum.add_value(l4_frame, true);
    icmp_checksum.add_bytes(l4_data, true);
    let icmp_checksum = icmp_checksum.finish();

    if icmp_checksum != 0 {
        eprintln!("ICMP checksum mismatch: {:#06x}", icmp_checksum);
        return false;
    }

    l4_data == expect_l4_data
}

#[allow(clippy::too_many_arguments)]
fn ping_once(
    socket: &mut RawSocket,
    poll: &mut PollChannel,
    timer: &mut TimerFd,
    info: &PingRouting,
    timeout: Duration,
    data_len: usize,
    id: u16,
    seq: u16,
) -> Result<bool, Error> {
    let mut buffer = [0; 4096];

    let source_ip = match info.source_ip {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => todo!(),
    };
    let destination_ip = match info.destination_ip {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => todo!(),
    };
    let mut l4_data = Vec::with_capacity(data_len);

    for _ in 0..data_len {
        l4_data.push(rand::random());
    }

    let ip_len = (size_of::<Ipv4Frame>() + size_of::<IcmpV4Frame>() + data_len)
        .try_into()
        .unwrap();

    let l2_frame = EthernetFrame {
        source_mac: (info.source_mac.octets()).into(),
        destination_mac: (info.gateway_mac.octets()).into(),
        ethertype: EtherType::IPV4.to_network_order(),
    };
    let mut l3_frame = Ipv4Frame {
        source_address: u32::from(source_ip).to_network_order(),
        destination_address: u32::from(destination_ip).to_network_order(),
        protocol: IpProtocol::ICMP,
        version_length: 0x45,
        total_length: u16::to_network_order(ip_len),
        flags_frag: u16::to_network_order(0x4000),
        id: u16::to_network_order(0),
        ttl: 255,
        ..Ipv4Frame::zeroed()
    };
    let mut l4_frame = IcmpV4Frame {
        ty: 8,
        code: 0,
        checksum: u16::to_network_order(0),
        rest: u32::to_network_order(((id as u32) << 16) | (seq as u32)),
    };

    let mut ip_checksum = InetChecksum::new();
    ip_checksum.add_value(&l3_frame, true);
    l3_frame.header_checksum = ip_checksum.finish().to_network_order();

    let mut icmp_checksum = InetChecksum::new();
    icmp_checksum.add_value(&l4_frame, true);
    icmp_checksum.add_bytes(&l4_data, true);
    l4_frame.checksum = icmp_checksum.finish().to_network_order();

    let mut packet = vec![];
    packet.extend_from_slice(bytemuck::bytes_of(&l2_frame));
    packet.extend_from_slice(bytemuck::bytes_of(&l3_frame));
    packet.extend_from_slice(bytemuck::bytes_of(&l4_frame));
    packet.extend_from_slice(&l4_data);

    timer.start(timeout)?;
    socket.send(&packet)?;

    loop {
        let (fd, result) = poll.wait(None, true)?.unwrap();
        result?;

        match fd {
            fd if fd == socket.as_raw_fd() => {
                // TODO
                let len = socket.recv(&mut buffer)?;
                if validate_ping_reply(&buffer[..len], source_ip, destination_ip, &l4_data, id, seq)
                {
                    return Ok(true);
                }
            }
            fd if fd == timer.as_raw_fd() => {
                return Ok(false);
            }
            _ => unreachable!(),
        }
    }
}

fn ping(
    input_address: &str,
    times: usize,
    data_len: usize,
    interval: Duration,
    timeout: Duration,
) -> Result<PingStats, Error> {
    let address = if let Ok(ip) = IpAddr::from_str(input_address) {
        ip
    } else {
        let socket_address = format!("{input_address}:0");
        let mut addrs = ToSocketAddrs::to_socket_addrs(&socket_address)?;
        let addr = addrs
            .next()
            .ok_or_else(|| Error::UnresolvedAddress(input_address.into()))?;
        addr.ip()
    };

    println!("Pinging {input_address:?} ({address})");

    let routing = resolve_routing(address)?;

    let mut stats = PingStats {
        packets_sent: 0,
        packets_received: 0,
    };
    let mut poll = PollChannel::new()?;
    let mut timer = TimerFd::new(false, false)?;
    let mut socket = RawSocket::bind(routing.interface_id)?;

    poll.add(timer.as_raw_fd())?;
    poll.add(socket.as_raw_fd())?;

    let id = rand::random();
    for i in 0..times {
        if INTERRUPTED.load(Ordering::Acquire) {
            break;
        }

        let result = ping_once(
            &mut socket,
            &mut poll,
            &mut timer,
            &routing,
            timeout,
            data_len,
            id,
            i as u16,
        )?;
        stats.packets_sent += 1;

        if result {
            stats.packets_received += 1;
            println!("[{}/{}] {}: PONG", i + 1, times, address);
        }

        std::thread::sleep(interval);
    }

    Ok(stats)
}

static INTERRUPTED: AtomicBool = AtomicBool::new(false);

fn main() -> ExitCode {
    // set_signal_handler(Signal::Interrupted, SignalHandler::Function(interrupt));

    let args = Args::parse();

    let stats = match ping(
        &args.address,
        args.count,
        args.data_size,
        Duration::from_millis(args.inteval.into()),
        Duration::from_millis(args.timeout.into()),
    ) {
        Ok(stats) => stats,
        Err(error) => {
            eprintln!("ping: {}", error);
            return ExitCode::FAILURE;
        }
    };

    let loss = (stats.packets_sent - stats.packets_received) * 100 / stats.packets_sent;
    println!(
        "{} sent, {} received, {}% loss",
        stats.packets_sent, stats.packets_received, loss
    );

    ExitCode::SUCCESS
}
