#![feature(yggdrasil_os, rustc_private)]
use std::{
    fs::File,
    io::{self, BufWriter, Write},
    os::yggdrasil::{
        io::net::raw_socket::RawSocket,
        signal::{set_signal_handler, SignalHandler},
    },
    path::{Path, PathBuf},
    process::ExitCode,
};

use bytemuck::{Pod, Zeroable};
use clap::Parser;
use ::yggdrasil_abi::net::protocols::EthernetFrame;
use runtime::abi as yggdrasil_abi;
use yggdrasil_abi::process::Signal;

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct PcapFileHeader {
    magic: u32,
    version_major: u16,
    version_minor: u16,
    _0: [u32; 2],
    snap_len: u32,
    link_type: u32,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct PcapRecordHeader {
    ts_seconds: u32,
    ts_nanoseconds: u32,
    captured_len: u32,
    original_len: u32,
}

#[derive(Parser, Debug)]
struct Args {
    #[clap(
        short,
        help = "Write captured packets to a PCAP file (Default: write to stdout)"
    )]
    output: Option<PathBuf>,
    #[clap(
        short,
        help = "When writing PCAP output to a file, also print packet information to stdout"
    )]
    noisy: bool,
    #[clap(help = "Interface to capture from")]
    interface: String,
}

struct PcapFile {
    writer: BufWriter<File>,
    header_written: bool,
}

impl PcapFile {
    const MAGIC: u32 = 0xA1B2C3D4;
    const VERSION_MAJOR: u16 = 2;
    const VERSION_MINOR: u16 = 0;

    fn open<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        File::create(path)
            .map(BufWriter::new)
            .map(Self::from_writer)
    }

    fn from_writer(writer: BufWriter<File>) -> Self {
        Self {
            writer,
            header_written: false,
        }
    }

    fn write(&mut self, payload: &[u8]) -> io::Result<()> {
        if !self.header_written {
            // Write the file header
            let header = PcapFileHeader {
                magic: Self::MAGIC.to_be(),
                version_minor: Self::VERSION_MINOR.to_be(),
                version_major: Self::VERSION_MAJOR.to_be(),
                _0: [0; 2],
                snap_len: 4096u32.to_be(),
                link_type: 0x1u32.to_be(),
            };
            self.writer.write_all(bytemuck::bytes_of(&header))?;
            self.header_written = true;
        }

        let record = PcapRecordHeader {
            ts_seconds: 0,
            ts_nanoseconds: 0,
            captured_len: (payload.len() as u32).to_be(),
            original_len: (payload.len() as u32).to_be(),
        };
        self.writer.write_all(bytemuck::bytes_of(&record))?;
        self.writer.write_all(payload)?;
        Ok(())
    }
}

enum Output {
    File(PcapFile),
    Stdout,
}

impl Output {
    fn open(path: Option<&Path>) -> io::Result<Self> {
        match path {
            Some(path) => PcapFile::open(path).map(Self::File),
            None => Ok(Self::Stdout),
        }
    }

    fn write(&mut self, noisy: bool, payload: &[u8]) -> io::Result<()> {
        match self {
            Self::File(file) => {
                if noisy {
                    describe_packet(payload);
                }

                file.write(payload)
            }
            Self::Stdout => {
                describe_packet(payload);
                Ok(())
            }
        }
    }
}

fn describe_packet(payload: &[u8]) {
    if payload.len() < size_of::<EthernetFrame>() {
        eprintln!("Runt packet of {} bytes", payload.len());
        return;
    }
    let ethernet: &EthernetFrame = bytemuck::from_bytes(&payload[..size_of::<EthernetFrame>()]);
    println!(
        "{} > {}, length {}",
        ethernet.source_mac,
        ethernet.destination_mac,
        payload.len() - size_of::<EthernetFrame>()
    );
}

fn run(args: &Args) -> io::Result<()> {
    let mut output = Output::open(args.output.as_deref())?;
    let socket = RawSocket::bind(args.interface.as_ref())?;
    let mut buffer = [0; 4096];

    // TODO packet timestamps
    loop {
        let len = socket.recv(&mut buffer)?;
        output.write(args.noisy, &buffer[..len])?;
    }
}

fn sigint_handler(_: Signal) {
    log::debug!("SIGINT handler!")
}

fn main() -> ExitCode {
    logsink::setup_logging(true);
    let _ = set_signal_handler(Signal::Interrupted, SignalHandler::Function(sigint_handler));

    let args = Args::parse();
    match run(&args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}
