use crate::packet::ip::IpNextLevelProtocol;
use pnet::packet::Packet;
use std::net::Ipv4Addr;

/// IPv4 Header Length
pub const IPV4_HEADER_LEN: usize = pnet::packet::ipv4::MutableIpv4Packet::minimum_packet_size();
/// IPv4 Header Bytes (32 bits)
pub const IPV4_HEADER_BYTES: usize = 4;

/// Represents the IPv4 flags.
#[derive(Clone, Debug, PartialEq)]
pub enum Ipv4Flags {
    Reserved,
    DontFragment,
    MoreFragments,
}

impl Ipv4Flags {
    pub fn from_u8(n: u8) -> Ipv4Flags {
        match n {
            0b010 => Ipv4Flags::DontFragment,
            0b001 => Ipv4Flags::MoreFragments,
            _ => Ipv4Flags::Reserved,
        }
    }
}

/// Represents the IPv4 options.
#[derive(Clone, Debug, PartialEq)]
pub enum Ipv4Option {
    EndOfOptionsList,
    NoOperation,
    Security,
    LooseSourceRoute,
    TimeStamp,
    ExtendedSecurity,
    CommercialSecurity,
    RecordRoute,
    StreamId,
    StrictSourceRoute,
    ExperimentalMeasurement,
    MtuProbe,
    MtuReply,
    ExperimentalFlowControl,
    ExperimentalAccessControl,
    Encode,
    ImiTrafficDescriptor,
    ExtendedInternetProtocol,
    Traceroute,
    AddressExtension,
    RouterAlert,
    SelectiveDirectedBroadcast,
    DynamicPacketState,
    UpstreamMulticastPacket,
    QuickStart,
    Rfc3692StyleExperiment,
    Unknown(u8),
}

impl Ipv4Option {
    /// Constructs a new Ipv4Option from u8
    pub fn from_u8(n: u8) -> Ipv4Option {
        match n {
            0 => Ipv4Option::EndOfOptionsList,
            1 => Ipv4Option::NoOperation,
            2 => Ipv4Option::Security,
            3 => Ipv4Option::LooseSourceRoute,
            4 => Ipv4Option::TimeStamp,
            5 => Ipv4Option::ExtendedSecurity,
            6 => Ipv4Option::CommercialSecurity,
            7 => Ipv4Option::RecordRoute,
            8 => Ipv4Option::StreamId,
            9 => Ipv4Option::StrictSourceRoute,
            10 => Ipv4Option::ExperimentalMeasurement,
            11 => Ipv4Option::MtuProbe,
            12 => Ipv4Option::MtuReply,
            13 => Ipv4Option::ExperimentalFlowControl,
            14 => Ipv4Option::ExperimentalAccessControl,
            15 => Ipv4Option::Encode,
            16 => Ipv4Option::ImiTrafficDescriptor,
            17 => Ipv4Option::ExtendedInternetProtocol,
            18 => Ipv4Option::Traceroute,
            19 => Ipv4Option::AddressExtension,
            20 => Ipv4Option::RouterAlert,
            21 => Ipv4Option::SelectiveDirectedBroadcast,
            23 => Ipv4Option::DynamicPacketState,
            24 => Ipv4Option::UpstreamMulticastPacket,
            25 => Ipv4Option::QuickStart,
            30 => Ipv4Option::Rfc3692StyleExperiment,
            _ => Ipv4Option::Unknown(n),
        }
    }
    /// Returns the number of the Ipv4Option
    pub fn number(&self) -> u8 {
        match *self {
            Ipv4Option::EndOfOptionsList => 0,
            Ipv4Option::NoOperation => 1,
            Ipv4Option::Security => 2,
            Ipv4Option::LooseSourceRoute => 3,
            Ipv4Option::TimeStamp => 4,
            Ipv4Option::ExtendedSecurity => 5,
            Ipv4Option::CommercialSecurity => 6,
            Ipv4Option::RecordRoute => 7,
            Ipv4Option::StreamId => 8,
            Ipv4Option::StrictSourceRoute => 9,
            Ipv4Option::ExperimentalMeasurement => 10,
            Ipv4Option::MtuProbe => 11,
            Ipv4Option::MtuReply => 12,
            Ipv4Option::ExperimentalFlowControl => 13,
            Ipv4Option::ExperimentalAccessControl => 14,
            Ipv4Option::Encode => 15,
            Ipv4Option::ImiTrafficDescriptor => 16,
            Ipv4Option::ExtendedInternetProtocol => 17,
            Ipv4Option::Traceroute => 18,
            Ipv4Option::AddressExtension => 19,
            Ipv4Option::RouterAlert => 20,
            Ipv4Option::SelectiveDirectedBroadcast => 21,
            Ipv4Option::DynamicPacketState => 23,
            Ipv4Option::UpstreamMulticastPacket => 24,
            Ipv4Option::QuickStart => 25,
            Ipv4Option::Rfc3692StyleExperiment => 30,
            Ipv4Option::Unknown(n) => n,
        }
    }
}

/// Represents an IPv4 packet.
#[derive(Clone, Debug, PartialEq)]
pub struct Ipv4Packet {
    /// IP version
    pub version: u8,
    /// IP header length
    pub header_length: u8,
    /// Differentiated Services Code Point (DSCP)
    pub dscp: u8,
    /// Explicit Congestion Notification (ECN)
    pub ecn: u8,
    /// Total length
    pub total_length: u16,
    /// Identification
    pub identification: u16,
    /// Flags
    pub flags: u8,
    /// Fragment offset
    pub fragment_offset: u16,
    /// Time to live
    pub ttl: u8,
    /// Next level protocol
    pub next_level_protocol: IpNextLevelProtocol,
    /// Checksum
    pub checksum: u16,
    /// Source IP address
    pub source: Ipv4Addr,
    /// Destination IP address
    pub destination: Ipv4Addr,
    /// Options
    pub options: Vec<Ipv4Option>,
    /// Payload
    pub payload: Vec<u8>,
}

impl Ipv4Packet {
    pub(crate) fn from_pnet_packet(packet: &pnet::packet::ipv4::Ipv4Packet) -> Ipv4Packet {
        Ipv4Packet {
            version: packet.get_version(),
            header_length: packet.get_header_length(),
            dscp: packet.get_dscp(),
            ecn: packet.get_ecn(),
            total_length: packet.get_total_length(),
            identification: packet.get_identification(),
            flags: packet.get_flags(),
            fragment_offset: packet.get_fragment_offset(),
            ttl: packet.get_ttl(),
            next_level_protocol: IpNextLevelProtocol::from_u8(packet.get_next_level_protocol().0),
            checksum: packet.get_checksum(),
            source: packet.get_source(),
            destination: packet.get_destination(),
            options: packet
                .get_options_iter()
                .map(|opt| Ipv4Option::from_u8(opt.get_number().0))
                .collect(),
            payload: packet.payload().to_vec(),
        }
    }
    /// Constructs a new Ipv4Packet from bytes
    pub fn from_bytes(packet: &[u8]) -> Ipv4Packet {
        let ipv4_packet = pnet::packet::ipv4::Ipv4Packet::new(packet).unwrap();
        Ipv4Packet::from_pnet_packet(&ipv4_packet)
    }
}

/// Build IPv4 packet
pub(crate) fn build_ipv4_packet(
    ipv4_packet: &mut pnet::packet::ipv4::MutableIpv4Packet,
    src_ip: Ipv4Addr,
    dst_ip: Ipv4Addr,
    next_protocol: IpNextLevelProtocol,
) {
    ipv4_packet.set_header_length((IPV4_HEADER_LEN / IPV4_HEADER_BYTES) as u8);
    ipv4_packet.set_source(src_ip);
    ipv4_packet.set_destination(dst_ip);
    ipv4_packet.set_identification(rand::random::<u16>());
    ipv4_packet.set_ttl(64);
    ipv4_packet.set_version(4);
    ipv4_packet.set_flags(pnet::packet::ipv4::Ipv4Flags::DontFragment);
    match next_protocol {
        IpNextLevelProtocol::Tcp => {
            ipv4_packet.set_total_length(52);
            ipv4_packet.set_next_level_protocol(pnet::packet::ip::IpNextHeaderProtocols::Tcp);
        }
        IpNextLevelProtocol::Udp => {
            ipv4_packet.set_total_length(28);
            ipv4_packet.set_next_level_protocol(pnet::packet::ip::IpNextHeaderProtocols::Udp);
        }
        IpNextLevelProtocol::Icmp => {
            ipv4_packet.set_total_length(28);
            ipv4_packet.set_next_level_protocol(pnet::packet::ip::IpNextHeaderProtocols::Icmp);
        }
        _ => {}
    }
    let checksum = pnet::packet::ipv4::checksum(&ipv4_packet.to_immutable());
    ipv4_packet.set_checksum(checksum);
}

/// IPv4 Packet Builder
#[derive(Clone, Debug)]
pub struct Ipv4PacketBuilder {
    pub src_ip: Ipv4Addr,
    pub dst_ip: Ipv4Addr,
    pub next_protocol: IpNextLevelProtocol,
    pub total_length: Option<u16>,
    pub identification: Option<u16>,
    pub ttl: Option<u8>,
    pub flags: Option<Ipv4Flags>,
}

impl Ipv4PacketBuilder {
    /// Constructs a new Ipv4PacketBuilder
    pub fn new(src_ip: Ipv4Addr, dst_ip: Ipv4Addr, next_protocol: IpNextLevelProtocol) -> Self {
        Ipv4PacketBuilder {
            src_ip,
            dst_ip,
            next_protocol,
            total_length: None,
            identification: None,
            ttl: None,
            flags: None,
        }
    }
    /// Builds IPv4 packet and return bytes
    pub fn build(&self) -> Vec<u8> {
        let mut buffer = vec![0; IPV4_HEADER_LEN];
        let mut ipv4_packet = pnet::packet::ipv4::MutableIpv4Packet::new(&mut buffer).unwrap();
        ipv4_packet.set_header_length((IPV4_HEADER_LEN / IPV4_HEADER_BYTES) as u8);
        ipv4_packet.set_source(self.src_ip);
        ipv4_packet.set_destination(self.dst_ip);
        ipv4_packet.set_identification(self.identification.unwrap_or(rand::random::<u16>()));
        ipv4_packet.set_ttl(self.ttl.unwrap_or(64));
        ipv4_packet.set_version(4);
        if let Some(flags) = &self.flags {
            match flags {
                Ipv4Flags::DontFragment => {
                    ipv4_packet.set_flags(pnet::packet::ipv4::Ipv4Flags::DontFragment);
                }
                Ipv4Flags::MoreFragments => {
                    ipv4_packet.set_flags(pnet::packet::ipv4::Ipv4Flags::MoreFragments);
                }
                _ => {}
            }
        } else {
            ipv4_packet.set_flags(pnet::packet::ipv4::Ipv4Flags::DontFragment);
        }
        match self.next_protocol {
            IpNextLevelProtocol::Tcp => {
                if let Some(total_length) = self.total_length {
                    ipv4_packet.set_total_length(total_length);
                } else {
                    ipv4_packet.set_total_length(52);
                }
                ipv4_packet.set_next_level_protocol(pnet::packet::ip::IpNextHeaderProtocols::Tcp);
            }
            IpNextLevelProtocol::Udp => {
                if let Some(total_length) = self.total_length {
                    ipv4_packet.set_total_length(total_length);
                } else {
                    ipv4_packet.set_total_length(28);
                }
                ipv4_packet.set_next_level_protocol(pnet::packet::ip::IpNextHeaderProtocols::Udp);
            }
            IpNextLevelProtocol::Icmp => {
                if let Some(total_length) = self.total_length {
                    ipv4_packet.set_total_length(total_length);
                } else {
                    ipv4_packet.set_total_length(28);
                }
                ipv4_packet.set_next_level_protocol(pnet::packet::ip::IpNextHeaderProtocols::Icmp);
            }
            _ => {}
        }
        let checksum = pnet::packet::ipv4::checksum(&ipv4_packet.to_immutable());
        ipv4_packet.set_checksum(checksum);
        ipv4_packet.packet().to_vec()
    }
}