use std::net::Ipv4Addr;

/// Overhead breakdown for VPN tunnel encapsulation.
#[derive(Debug, Clone)]
pub struct TunnelOverhead {
    /// Outer IP header: 20 bytes (IPv4, no options).
    pub ip_header: u16,
    /// TCP header: typically 32 bytes (20 base + 12 for timestamps).
    pub tcp_header: u16,
    /// WebSocket framing: ~6 bytes (2 base + 4 mask from client).
    pub ws_framing: u16,
    /// VPN binary frame header: 5 bytes [type:1B][length:4B].
    pub vpn_header: u16,
    /// Noise AEAD tag: 16 bytes (Poly1305).
    pub noise_tag: u16,
}

impl TunnelOverhead {
    /// Conservative default overhead estimate.
    pub fn default_overhead() -> Self {
        Self {
            ip_header: 20,
            tcp_header: 32,
            ws_framing: 6,
            vpn_header: 5,
            noise_tag: 16,
        }
    }

    /// Total encapsulation overhead in bytes.
    pub fn total(&self) -> u16 {
        self.ip_header + self.tcp_header + self.ws_framing + self.vpn_header + self.noise_tag
    }

    /// Compute effective TUN MTU given the underlying link MTU.
    pub fn effective_tun_mtu(&self, link_mtu: u16) -> u16 {
        link_mtu.saturating_sub(self.total())
    }
}

/// MTU configuration for the VPN tunnel.
#[derive(Debug, Clone)]
pub struct MtuConfig {
    /// Underlying link MTU (typically 1500 for Ethernet).
    pub link_mtu: u16,
    /// Computed effective TUN MTU.
    pub effective_mtu: u16,
    /// Whether to generate ICMP too-big for oversized packets.
    pub send_icmp_too_big: bool,
    /// Counter: oversized packets encountered.
    pub oversized_packets: u64,
    /// Counter: ICMP too-big messages generated.
    pub icmp_too_big_sent: u64,
}

impl MtuConfig {
    /// Create a new MTU config from the underlying link MTU.
    pub fn new(link_mtu: u16) -> Self {
        let overhead = TunnelOverhead::default_overhead();
        let effective = overhead.effective_tun_mtu(link_mtu);
        Self {
            link_mtu,
            effective_mtu: effective,
            send_icmp_too_big: true,
            oversized_packets: 0,
            icmp_too_big_sent: 0,
        }
    }

    /// Check if a packet exceeds the effective MTU.
    pub fn is_oversized(&self, packet_len: usize) -> bool {
        packet_len > self.effective_mtu as usize
    }
}

/// Action to take after checking MTU.
pub enum MtuAction {
    /// Packet is within MTU, forward normally.
    Forward,
    /// Packet is oversized; contains the ICMP too-big message to write back into TUN.
    SendIcmpTooBig(Vec<u8>),
}

/// Check packet against MTU config and return the appropriate action.
pub fn check_mtu(packet: &[u8], config: &MtuConfig) -> MtuAction {
    if !config.is_oversized(packet.len()) {
        return MtuAction::Forward;
    }

    if !config.send_icmp_too_big {
        return MtuAction::Forward;
    }

    match generate_icmp_too_big(packet, config.effective_mtu) {
        Some(icmp) => MtuAction::SendIcmpTooBig(icmp),
        None => MtuAction::Forward,
    }
}

/// Generate an ICMPv4 Destination Unreachable / Fragmentation Needed message.
///
/// Per RFC 792: Type 3, Code 4, with next-hop MTU in bytes 6-7 (RFC 1191).
/// Returns the complete IP + ICMP packet to write back into the TUN device.
pub fn generate_icmp_too_big(original_packet: &[u8], next_hop_mtu: u16) -> Option<Vec<u8>> {
    // Need at least 20 bytes of original IP header
    if original_packet.len() < 20 {
        return None;
    }

    // Verify it's IPv4
    if original_packet[0] >> 4 != 4 {
        return None;
    }

    // Parse source/dest from original IP header
    let src_ip = Ipv4Addr::new(
        original_packet[12],
        original_packet[13],
        original_packet[14],
        original_packet[15],
    );
    let dst_ip = Ipv4Addr::new(
        original_packet[16],
        original_packet[17],
        original_packet[18],
        original_packet[19],
    );

    // ICMP payload: IP header + first 8 bytes of original datagram (per RFC 792)
    let icmp_data_len = original_packet.len().min(28); // 20 IP header + 8 bytes
    let icmp_payload = &original_packet[..icmp_data_len];

    // Build ICMP message: type(1) + code(1) + checksum(2) + unused(2) + next_hop_mtu(2) + data
    let mut icmp = Vec::with_capacity(8 + icmp_data_len);
    icmp.push(3); // Type: Destination Unreachable
    icmp.push(4); // Code: Fragmentation Needed and DF was Set
    icmp.push(0); // Checksum placeholder
    icmp.push(0);
    icmp.push(0); // Unused
    icmp.push(0);
    icmp.extend_from_slice(&next_hop_mtu.to_be_bytes());
    icmp.extend_from_slice(icmp_payload);

    // Compute ICMP checksum
    let cksum = internet_checksum(&icmp);
    icmp[2] = (cksum >> 8) as u8;
    icmp[3] = (cksum & 0xff) as u8;

    // Build IP header (ICMP response: FROM tunnel gateway TO original source)
    let total_len = (20 + icmp.len()) as u16;
    let mut ip = Vec::with_capacity(total_len as usize);
    ip.push(0x45); // Version 4, IHL 5
    ip.push(0x00); // DSCP/ECN
    ip.extend_from_slice(&total_len.to_be_bytes());
    ip.extend_from_slice(&[0, 0]); // Identification
    ip.extend_from_slice(&[0x40, 0x00]); // Flags: Don't Fragment, Fragment Offset: 0
    ip.push(64); // TTL
    ip.push(1); // Protocol: ICMP
    ip.extend_from_slice(&[0, 0]); // Header checksum placeholder
    ip.extend_from_slice(&dst_ip.octets()); // Source: tunnel endpoint (was dst)
    ip.extend_from_slice(&src_ip.octets()); // Destination: original source

    // Compute IP header checksum
    let ip_cksum = internet_checksum(&ip[..20]);
    ip[10] = (ip_cksum >> 8) as u8;
    ip[11] = (ip_cksum & 0xff) as u8;

    ip.extend_from_slice(&icmp);
    Some(ip)
}

/// Standard Internet checksum (RFC 1071).
fn internet_checksum(data: &[u8]) -> u16 {
    let mut sum: u32 = 0;
    let mut i = 0;
    while i + 1 < data.len() {
        sum += u16::from_be_bytes([data[i], data[i + 1]]) as u32;
        i += 2;
    }
    if i < data.len() {
        sum += (data[i] as u32) << 8;
    }
    while sum >> 16 != 0 {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }
    !sum as u16
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn default_overhead_total() {
        let oh = TunnelOverhead::default_overhead();
        assert_eq!(oh.total(), 79); // 20+32+6+5+16
    }

    #[test]
    fn effective_mtu_for_ethernet() {
        let oh = TunnelOverhead::default_overhead();
        let mtu = oh.effective_tun_mtu(1500);
        assert_eq!(mtu, 1421); // 1500 - 79
    }

    #[test]
    fn effective_mtu_saturates_at_zero() {
        let oh = TunnelOverhead::default_overhead();
        let mtu = oh.effective_tun_mtu(50); // Less than overhead
        assert_eq!(mtu, 0);
    }

    #[test]
    fn mtu_config_default() {
        let config = MtuConfig::new(1500);
        assert_eq!(config.effective_mtu, 1421);
        assert_eq!(config.link_mtu, 1500);
        assert!(config.send_icmp_too_big);
    }

    #[test]
    fn is_oversized() {
        let config = MtuConfig::new(1500);
        assert!(!config.is_oversized(1421));
        assert!(config.is_oversized(1422));
    }

    #[test]
    fn icmp_too_big_generation() {
        // Craft a minimal IPv4 packet
        let mut original = vec![0u8; 28];
        original[0] = 0x45; // version 4, IHL 5
        original[2..4].copy_from_slice(&1500u16.to_be_bytes()); // total length
        original[9] = 6; // TCP
        original[12..16].copy_from_slice(&[10, 0, 0, 1]); // src IP
        original[16..20].copy_from_slice(&[10, 0, 0, 2]); // dst IP

        let icmp_pkt = generate_icmp_too_big(&original, 1421).unwrap();

        // Verify it's a valid IPv4 packet
        assert_eq!(icmp_pkt[0] >> 4, 4); // IPv4
        assert_eq!(icmp_pkt[9], 1); // ICMP protocol

        // Source should be original dst (10.0.0.2)
        assert_eq!(&icmp_pkt[12..16], &[10, 0, 0, 2]);
        // Destination should be original src (10.0.0.1)
        assert_eq!(&icmp_pkt[16..20], &[10, 0, 0, 1]);

        // ICMP type 3, code 4
        assert_eq!(icmp_pkt[20], 3);
        assert_eq!(icmp_pkt[21], 4);

        // Next-hop MTU at ICMP bytes 6-7 (offset 26-27 in IP packet)
        let mtu = u16::from_be_bytes([icmp_pkt[26], icmp_pkt[27]]);
        assert_eq!(mtu, 1421);
    }

    #[test]
    fn icmp_too_big_rejects_short_packet() {
        let short = vec![0u8; 10];
        assert!(generate_icmp_too_big(&short, 1421).is_none());
    }

    #[test]
    fn icmp_too_big_rejects_non_ipv4() {
        let mut pkt = vec![0u8; 40];
        pkt[0] = 0x60; // IPv6
        assert!(generate_icmp_too_big(&pkt, 1421).is_none());
    }

    #[test]
    fn icmp_checksum_valid() {
        let mut original = vec![0u8; 28];
        original[0] = 0x45;
        original[2..4].copy_from_slice(&1500u16.to_be_bytes());
        original[9] = 6;
        original[12..16].copy_from_slice(&[192, 168, 1, 100]);
        original[16..20].copy_from_slice(&[10, 8, 0, 1]);

        let icmp_pkt = generate_icmp_too_big(&original, 1420).unwrap();

        // Verify IP header checksum
        let ip_cksum = internet_checksum(&icmp_pkt[..20]);
        assert_eq!(ip_cksum, 0, "IP header checksum should verify to 0");

        // Verify ICMP checksum
        let icmp_cksum = internet_checksum(&icmp_pkt[20..]);
        assert_eq!(icmp_cksum, 0, "ICMP checksum should verify to 0");
    }

    #[test]
    fn check_mtu_forward() {
        let config = MtuConfig::new(1500);
        let pkt = vec![0u8; 1421]; // Exactly at MTU
        assert!(matches!(check_mtu(&pkt, &config), MtuAction::Forward));
    }

    #[test]
    fn check_mtu_oversized_generates_icmp() {
        let config = MtuConfig::new(1500);
        let mut pkt = vec![0u8; 1500];
        pkt[0] = 0x45; // Valid IPv4
        pkt[12..16].copy_from_slice(&[10, 0, 0, 1]);
        pkt[16..20].copy_from_slice(&[10, 0, 0, 2]);

        match check_mtu(&pkt, &config) {
            MtuAction::SendIcmpTooBig(icmp) => {
                assert_eq!(icmp[20], 3); // ICMP type
                assert_eq!(icmp[21], 4); // ICMP code
            }
            MtuAction::Forward => panic!("Expected SendIcmpTooBig"),
        }
    }
}