//! Raw AF_PACKET sender that constructs complete Ethernet + IP + UDP frames
//! in userspace and sends them directly to the NIC, bypassing the kernel's
//! UDP/IP stack entirely.

use std::fs;
use std::io::{self, BufRead, BufReader};
use std::net::{Ipv4Addr, ToSocketAddrs, UdpSocket};
use std::os::unix::io::AsRawFd;

/// Size of Ethernet + IPv4 + UDP headers combined.
const HEADER_SIZE: usize = 14 + 20 + 8; // 42 bytes

/// PACKET_QDISC_BYPASS socket option (skip the qdisc layer).
const PACKET_QDISC_BYPASS: libc::c_int = 20;

/// A raw AF_PACKET sender that constructs full Ethernet frames.
pub struct RawSender {
    fd: libc::c_int,
    ifindex: i32,
    /// Pre-built header template (Ethernet + IP + UDP).
    /// Only the IP total_length, IP checksum, and UDP length fields
    /// need to be patched per packet.
    header_template: [u8; HEADER_SIZE],
}

impl RawSender {
    /// Creates a new RawSender.
    ///
    /// Auto-discovers the network interface, MAC addresses, and source IP
    /// from the routing table and ARP cache.
    ///
    /// Requires `CAP_NET_RAW` or root.
    pub fn new(dst_host: &str, dst_port: u16, src_port: u16, interface: Option<&str>) -> io::Result<Self> {
        // Resolve destination IP.
        let dst_ip: Ipv4Addr = dst_host
            .parse()
            .or_else(|_| {
                (dst_host, dst_port)
                    .to_socket_addrs()?
                    .find_map(|a| match a {
                        std::net::SocketAddr::V4(v4) => Some(*v4.ip()),
                        _ => None,
                    })
                    .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "Could not resolve host"))
            })?;

        // Discover or use the specified source interface and IP.
        let (ifname, ifindex, src_ip) = if let Some(ifname) = interface {
            let ifindex = get_ifindex(ifname)?;
            let src_ip = get_interface_ip(ifname)?;
            (ifname.to_owned(), ifindex, src_ip)
        } else {
            let probe = UdpSocket::bind("0.0.0.0:0")?;
            probe.connect((dst_ip, dst_port))?;
            let src_ip: Ipv4Addr = match probe.local_addr()? {
                std::net::SocketAddr::V4(a) => *a.ip(),
                _ => return Err(io::Error::new(io::ErrorKind::Other, "Not IPv4")),
            };
            let (ifname, ifindex) = discover_interface(dst_ip)?;
            (ifname, ifindex, src_ip)
        };
        eprintln!("[raw] interface: {} (index {})", ifname, ifindex);

        // Read source MAC from sysfs.
        let src_mac = read_mac_from_sysfs(&ifname)?;
        eprintln!("[raw] src MAC: {:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}",
            src_mac[0], src_mac[1], src_mac[2], src_mac[3], src_mac[4], src_mac[5]);

        // Look up destination MAC from ARP cache.
        let dst_mac = lookup_arp(dst_ip, &ifname)?;
        eprintln!("[raw] dst MAC: {:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}",
            dst_mac[0], dst_mac[1], dst_mac[2], dst_mac[3], dst_mac[4], dst_mac[5]);

        eprintln!("[raw] {}:{} -> {}:{}", src_ip, src_port, dst_ip, dst_port);

        // Create AF_PACKET raw socket.
        let fd = unsafe {
            libc::socket(
                libc::AF_PACKET,
                libc::SOCK_RAW,
                (libc::ETH_P_IP as u16).to_be() as i32,
            )
        };
        if fd < 0 {
            return Err(io::Error::last_os_error());
        }

        // Bind to the interface.
        let mut sll: libc::sockaddr_ll = unsafe { std::mem::zeroed() };
        sll.sll_family = libc::AF_PACKET as u16;
        sll.sll_protocol = (libc::ETH_P_IP as u16).to_be();
        sll.sll_ifindex = ifindex;
        let ret = unsafe {
            libc::bind(
                fd,
                &sll as *const libc::sockaddr_ll as *const libc::sockaddr,
                std::mem::size_of::<libc::sockaddr_ll>() as libc::socklen_t,
            )
        };
        if ret < 0 {
            unsafe { libc::close(fd); }
            return Err(io::Error::last_os_error());
        }

        // Skip qdisc for lower latency and higher throughput.
        let bypass: libc::c_int = 1;
        unsafe {
            libc::setsockopt(
                fd,
                libc::SOL_PACKET,
                PACKET_QDISC_BYPASS,
                &bypass as *const libc::c_int as *const libc::c_void,
                std::mem::size_of::<libc::c_int>() as libc::socklen_t,
            );
        }

        // Increase send buffer.
        let sndbuf: libc::c_int = 25 * 1024 * 1024;
        unsafe {
            libc::setsockopt(
                fd,
                libc::SOL_SOCKET,
                libc::SO_SNDBUF,
                &sndbuf as *const libc::c_int as *const libc::c_void,
                std::mem::size_of::<libc::c_int>() as libc::socklen_t,
            );
        }

        // Build the header template.
        let header_template = build_header_template(
            &src_mac, &dst_mac, src_ip, dst_ip, src_port, dst_port,
        );

        Ok(RawSender {
            fd,
            ifindex,
            header_template,
        })
    }

    /// Sends a batch of payloads as raw Ethernet frames using sendmmsg.
    ///
    /// `frames` is a pre-allocated buffer of frame data.
    /// Each frame is HEADER_SIZE + payload bytes.
    /// Returns the number of frames successfully sent.
    pub fn send_batch(&self, payloads: &[&[u8]], frame_buf: &mut Vec<u8>) -> io::Result<usize> {
        let count = payloads.len();
        if count == 0 {
            return Ok(0);
        }

        // Build all frames into the contiguous buffer.
        frame_buf.clear();
        let mut offsets: Vec<(usize, usize)> = Vec::with_capacity(count);

        for payload in payloads {
            let frame_start = frame_buf.len();
            let total_len = (20 + 8 + payload.len()) as u16; // IP total length
            let udp_len = (8 + payload.len()) as u16;

            // Copy header template.
            frame_buf.extend_from_slice(&self.header_template);

            // Patch IP total length (bytes 16-17 of the frame = offset 2-3 in IP header).
            let ip_start = frame_start + 14;
            frame_buf[ip_start + 2] = (total_len >> 8) as u8;
            frame_buf[ip_start + 3] = total_len as u8;

            // Patch IP header checksum (bytes 10-11 of IP header).
            // Zero the checksum field first, then compute.
            frame_buf[ip_start + 10] = 0;
            frame_buf[ip_start + 11] = 0;
            let cksum = ip_checksum(&frame_buf[ip_start..ip_start + 20]);
            frame_buf[ip_start + 10] = (cksum >> 8) as u8;
            frame_buf[ip_start + 11] = cksum as u8;

            // Patch UDP length (bytes 4-5 of UDP header).
            let udp_start = frame_start + 14 + 20;
            frame_buf[udp_start + 4] = (udp_len >> 8) as u8;
            frame_buf[udp_start + 5] = udp_len as u8;

            // UDP checksum = 0 (optional for IPv4).
            frame_buf[udp_start + 6] = 0;
            frame_buf[udp_start + 7] = 0;

            // Append payload.
            frame_buf.extend_from_slice(payload);

            let frame_end = frame_buf.len();
            offsets.push((frame_start, frame_end));
        }

        // Build iovec and mmsghdr arrays for sendmmsg.
        let mut iovecs: Vec<libc::iovec> = Vec::with_capacity(count);
        let mut msghdrs: Vec<libc::mmsghdr> = Vec::with_capacity(count);

        // Destination sockaddr_ll.
        let mut sll: libc::sockaddr_ll = unsafe { std::mem::zeroed() };
        sll.sll_family = libc::AF_PACKET as u16;
        sll.sll_protocol = (libc::ETH_P_IP as u16).to_be();
        sll.sll_ifindex = self.ifindex;
        sll.sll_halen = 6;
        // dst MAC is in the first 6 bytes of each frame, but sockaddr_ll
        // also needs it. Copy from the template.
        sll.sll_addr[..6].copy_from_slice(&self.header_template[..6]);

        for &(start, end) in &offsets {
            iovecs.push(libc::iovec {
                iov_base: frame_buf[start..end].as_ptr() as *mut libc::c_void,
                iov_len: end - start,
            });
        }

        for iov in &iovecs {
            let mut mhdr: libc::mmsghdr = unsafe { std::mem::zeroed() };
            mhdr.msg_hdr.msg_iov = iov as *const libc::iovec as *mut libc::iovec;
            mhdr.msg_hdr.msg_iovlen = 1;
            mhdr.msg_hdr.msg_name = &sll as *const libc::sockaddr_ll as *mut libc::c_void;
            mhdr.msg_hdr.msg_namelen = std::mem::size_of::<libc::sockaddr_ll>() as u32;
            msghdrs.push(mhdr);
        }

        let ret = unsafe {
            libc::sendmmsg(
                self.fd,
                msghdrs.as_mut_ptr(),
                count as libc::c_uint,
                0,
            )
        };

        if ret < 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(ret as usize)
        }
    }

    /// Returns the header size so callers can compute total frame sizes.
    pub fn header_size(&self) -> usize {
        HEADER_SIZE
    }
}

impl Drop for RawSender {
    fn drop(&mut self) {
        unsafe {
            libc::close(self.fd);
        }
    }
}

/// Builds the 42-byte header template: Ethernet (14) + IP (20) + UDP (8).
///
/// Fields that vary per packet (IP total_length, IP checksum, UDP length)
/// are filled with placeholder values and must be patched before sending.
fn build_header_template(
    src_mac: &[u8; 6],
    dst_mac: &[u8; 6],
    src_ip: Ipv4Addr,
    dst_ip: Ipv4Addr,
    src_port: u16,
    dst_port: u16,
) -> [u8; HEADER_SIZE] {
    let mut h = [0u8; HEADER_SIZE];

    // -- Ethernet header (14 bytes) --
    h[0..6].copy_from_slice(dst_mac);
    h[6..12].copy_from_slice(src_mac);
    h[12] = 0x08; // EtherType: IPv4
    h[13] = 0x00;

    // -- IPv4 header (20 bytes, no options) --
    let ip = &mut h[14..34];
    ip[0] = 0x45; // Version 4, IHL 5 (20 bytes)
    ip[1] = 0x00; // DSCP/ECN
    // ip[2..4] = total length (patched per packet)
    ip[4] = 0x00; // Identification
    ip[5] = 0x00;
    ip[6] = 0x40; // Flags: Don't Fragment
    ip[7] = 0x00; // Fragment offset
    ip[8] = 64; // TTL
    ip[9] = 17; // Protocol: UDP
    // ip[10..12] = header checksum (patched per packet)
    ip[12..16].copy_from_slice(&src_ip.octets());
    ip[16..20].copy_from_slice(&dst_ip.octets());

    // -- UDP header (8 bytes) --
    let udp = &mut h[34..42];
    udp[0] = (src_port >> 8) as u8;
    udp[1] = src_port as u8;
    udp[2] = (dst_port >> 8) as u8;
    udp[3] = dst_port as u8;
    // udp[4..6] = length (patched per packet)
    // udp[6..8] = checksum (set to 0, optional for IPv4)

    h
}

/// Computes the IPv4 header checksum (one's complement of the one's
/// complement sum of all 16-bit words in the header).
fn ip_checksum(header: &[u8]) -> u16 {
    let mut sum: u32 = 0;
    for i in (0..header.len()).step_by(2) {
        let word = if i + 1 < header.len() {
            ((header[i] as u32) << 8) | (header[i + 1] as u32)
        } else {
            (header[i] as u32) << 8
        };
        sum += word;
    }
    // Fold carry bits.
    while sum > 0xFFFF {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }
    !(sum as u16)
}

/// Gets the ifindex for a named interface from sysfs.
fn get_ifindex(ifname: &str) -> io::Result<i32> {
    let path = format!("/sys/class/net/{}/ifindex", ifname);
    let s = fs::read_to_string(&path).map_err(|e| {
        io::Error::new(e.kind(), format!("Interface '{}' not found: {}", ifname, e))
    })?;
    s.trim().parse().map_err(|e| {
        io::Error::new(io::ErrorKind::InvalidData, format!("Bad ifindex: {}", e))
    })
}

/// Gets the IPv4 address of a named interface using ioctl.
fn get_interface_ip(ifname: &str) -> io::Result<Ipv4Addr> {
    let sock = UdpSocket::bind("0.0.0.0:0")?;
    let fd = sock.as_raw_fd();
    let mut ifr: libc::ifreq = unsafe { std::mem::zeroed() };
    let name_bytes = ifname.as_bytes();
    let copy_len = name_bytes.len().min(libc::IFNAMSIZ - 1);
    unsafe {
        std::ptr::copy_nonoverlapping(
            name_bytes.as_ptr(),
            ifr.ifr_name.as_mut_ptr() as *mut u8,
            copy_len,
        );
    }
    let ret = unsafe { libc::ioctl(fd, libc::SIOCGIFADDR, &mut ifr) };
    if ret < 0 {
        return Err(io::Error::new(
            io::ErrorKind::Other,
            format!("Could not get IP for interface '{}'", ifname),
        ));
    }
    let addr = unsafe { &*(&ifr.ifr_ifru as *const _ as *const libc::sockaddr_in) };
    Ok(Ipv4Addr::from(u32::from_be(addr.sin_addr.s_addr)))
}

/// Discovers the outbound interface name and index for a destination IP.
fn discover_interface(dst_ip: Ipv4Addr) -> io::Result<(String, i32)> {
    // Use a temporary UDP socket to discover the interface.
    let sock = UdpSocket::bind("0.0.0.0:0")?;
    sock.connect((dst_ip, 80))?;

    // Get the interface index via SO_BINDTODEVICE or by reading /proc/net/route.
    // Simpler: use the socket's bound address to find the matching interface.
    let local_ip = match sock.local_addr()? {
        std::net::SocketAddr::V4(a) => *a.ip(),
        _ => return Err(io::Error::new(io::ErrorKind::Other, "Not IPv4")),
    };

    // Scan /proc/net/if_inet6 or /proc/net/fib_trie... actually, simplest
    // is to iterate /sys/class/net/*/address and check IPs.
    // But even simpler: parse `ip route get` output... that's ugly.
    //
    // Let's iterate interfaces and match by IP.
    let entries = fs::read_dir("/sys/class/net/")?;
    for entry in entries {
        let entry = entry?;
        let ifname = entry.file_name().to_string_lossy().into_owned();

        // Get the interface index.
        let ifindex_path = format!("/sys/class/net/{}/ifindex", ifname);
        let ifindex_str = match fs::read_to_string(&ifindex_path) {
            Ok(s) => s,
            Err(_) => continue,
        };
        let ifindex: i32 = match ifindex_str.trim().parse() {
            Ok(i) => i,
            Err(_) => continue,
        };

        // Check if this interface has the local IP using ioctl.
        let fd = sock.as_raw_fd();
        let mut ifr: libc::ifreq = unsafe { std::mem::zeroed() };
        let name_bytes = ifname.as_bytes();
        let copy_len = name_bytes.len().min(libc::IFNAMSIZ - 1);
        unsafe {
            std::ptr::copy_nonoverlapping(
                name_bytes.as_ptr(),
                ifr.ifr_name.as_mut_ptr() as *mut u8,
                copy_len,
            );
        }

        let ret = unsafe {
            libc::ioctl(fd, libc::SIOCGIFADDR, &mut ifr)
        };
        if ret < 0 {
            continue;
        }

        let addr = unsafe { &*(&ifr.ifr_ifru as *const _ as *const libc::sockaddr_in) };
        let iface_ip = Ipv4Addr::from(u32::from_be(addr.sin_addr.s_addr));

        if iface_ip == local_ip {
            return Ok((ifname, ifindex));
        }
    }

    Err(io::Error::new(
        io::ErrorKind::NotFound,
        format!("Could not find interface for IP {}", local_ip),
    ))
}

/// Reads the MAC address of a network interface from sysfs.
fn read_mac_from_sysfs(ifname: &str) -> io::Result<[u8; 6]> {
    let path = format!("/sys/class/net/{}/address", ifname);
    let mac_str = fs::read_to_string(path)?;
    parse_mac(mac_str.trim())
}

/// Looks up the MAC address for an IP in the kernel ARP cache.
fn lookup_arp(ip: Ipv4Addr, ifname: &str) -> io::Result<[u8; 6]> {
    let file = fs::File::open("/proc/net/arp")?;
    let reader = BufReader::new(file);

    let ip_str = ip.to_string();

    for line in reader.lines().skip(1) {
        let line = line?;
        let fields: Vec<&str> = line.split_whitespace().collect();
        if fields.len() >= 6 && fields[0] == ip_str && fields[5] == ifname {
            return parse_mac(fields[3]);
        }
    }

    Err(io::Error::new(
        io::ErrorKind::NotFound,
        format!("No ARP entry for {} on {}", ip, ifname),
    ))
}

/// Parses a MAC address string like "50:6b:4b:c3:fa:9c" into 6 bytes.
fn parse_mac(s: &str) -> io::Result<[u8; 6]> {
    let parts: Vec<&str> = s.split(':').collect();
    if parts.len() != 6 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("Invalid MAC: {}", s),
        ));
    }
    let mut mac = [0u8; 6];
    for (i, part) in parts.iter().enumerate() {
        mac[i] = u8::from_str_radix(part, 16).map_err(|_| {
            io::Error::new(io::ErrorKind::InvalidData, format!("Invalid MAC byte: {}", part))
        })?;
    }
    Ok(mac)
}