// Copyright (C) 2023 Andrew Rioux
// 
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
// 
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use crate::error;
use std::{
    net::Ipv4Addr,
    sync::atomic::{AtomicU16, Ordering},
};

pub struct EthernetPkt<'a> {
    pub(crate) data: &'a [u8],
}

impl<'a> EthernetPkt<'a> {
    pub fn destination_address(&self) -> &[u8; 6] {
        self.data[0..6].try_into().unwrap()
    }

    pub fn source_address(&self) -> &[u8; 6] {
        self.data[6..12].try_into().unwrap()
    }

    pub fn ether_type(&self) -> u16 {
        u16::from_be_bytes(self.data[12..14].try_into().unwrap())
    }

    pub fn get_layer3_pkt(&'_ self) -> error::Result<Layer3Pkt<'_>> {
        if self.data.len() < 14 {
            return Err(error::Error::PacketLengthInvalid);
        }

        match self.ether_type() {
            0x0800 => {
                if self.data.len() < 34 {
                    return Err(error::Error::PacketLengthInvalid);
                }

                let pkt = IPv4Pkt {
                    data: &self.data[14..],
                };

                if self.data.len() < (14 + pkt.total_length()).into() {
                    return Err(error::Error::PacketLengthInvalid);
                }

                Ok(Layer3Pkt::IPv4Pkt(pkt))
            }
            p => Err(error::Error::UnknownPacketType(p)),
        }
    }

    pub fn to_owned(&self) -> EthernetPacket {
        EthernetPacket {
            data: self.data.to_vec(),
        }
    }
}

pub enum Layer3Pkt<'a> {
    IPv4Pkt(IPv4Pkt<'a>),
}

pub struct IPv4Pkt<'a> {
    data: &'a [u8],
}

impl<'a> IPv4Pkt<'a> {
    pub fn version(&self) -> u8 {
        4
    }

    pub fn header_len(&self) -> u8 {
        4 * (0x0F & self.data[0])
    }

    pub fn type_of_service(&self) -> u8 {
        self.data[1] >> 2
    }

    pub fn ecn(&self) -> u8 {
        self.data[1] & 0x03
    }

    pub fn total_length(&self) -> u16 {
        u16::from_be_bytes(self.data[2..4].try_into().unwrap())
    }

    pub fn id(&self) -> u16 {
        u16::from_be_bytes(self.data[4..6].try_into().unwrap())
    }

    pub fn dont_fragment(&self) -> bool {
        self.data[6] & 0x40 != 0
    }

    pub fn more_fragments(&self) -> bool {
        self.data[6] & 0x80 != 0
    }

    pub fn fragment_offset(&self) -> u16 {
        u16::from_be_bytes([self.data[6] & 0x1F, self.data[7]]) * 8
    }

    pub fn ttl(&self) -> u8 {
        self.data[8]
    }

    pub fn protocol(&self) -> u8 {
        self.data[9]
    }

    pub fn header_checksum(&self) -> u16 {
        u16::from_be_bytes(self.data[10..12].try_into().unwrap())
    }

    pub fn source_ip(&self) -> Ipv4Addr {
        <&[u8] as TryInto<[u8; 4]>>::try_into(&self.data[12..16])
            .unwrap()
            .into()
    }

    pub fn dest_ip(&self) -> Ipv4Addr {
        <&[u8] as TryInto<[u8; 4]>>::try_into(&self.data[16..20])
            .unwrap()
            .into()
    }

    pub fn to_owned(&self) -> IPv4Packet {
        IPv4Packet {
            data: self.data.to_vec(),
        }
    }

    pub fn get_layer4_packet(&self) -> error::Result<Layer4Pkt<'a>> {
        match self.protocol() {
            17 => Ok(Layer4Pkt::UDP(UDPPkt {
                data: &self.data[(self.header_len() as usize)..],
            })),
            p => Err(error::Error::UnknownPacketType(p.into())),
        }
    }
}

pub enum Layer4Pkt<'a> {
    UDP(UDPPkt<'a>),
}

impl<'a> Layer4Pkt<'a> {
    pub fn len(&self) -> u16 {
        match self {
            Layer4Pkt::UDP(pkt) => pkt.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

pub struct UDPPkt<'a> {
    data: &'a [u8],
}

impl<'a> UDPPkt<'a> {
    pub fn srcport(&self) -> u16 {
        u16::from_be_bytes(self.data[0..2].try_into().unwrap())
    }

    pub fn dstport(&self) -> u16 {
        u16::from_be_bytes(self.data[2..4].try_into().unwrap())
    }

    pub fn len(&self) -> u16 {
        u16::from_be_bytes(self.data[4..6].try_into().unwrap())
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn to_owned(&self) -> UDPPacket {
        UDPPacket {
            data: self.data.to_vec(),
        }
    }

    pub fn get_data(&'_ self) -> &'_ [u8] {
        &self.data[8..(self.len() as usize)]
    }
}

#[derive(Debug, Clone)]
pub struct EthernetPacket {
    data: Vec<u8>,
}

impl EthernetPacket {
    pub fn construct(src: [u8; 6], dst: [u8; 6], packet: &Layer3Packet) -> EthernetPacket {
        match packet {
            Layer3Packet::IPv4(pkt) => EthernetPacket {
                data: [&dst[..], &src[..], &[0x08_u8, 0x00_u8][..], &*pkt.data].concat(),
            },
        }
    }

    #[inline]
    pub fn pkt(&'_ self) -> EthernetPkt<'_> {
        EthernetPkt { data: &self.data }
    }
}

#[derive(Clone)]
pub enum Layer3Packet {
    IPv4(IPv4Packet),
}

static IPV4_ID: AtomicU16 = AtomicU16::new(0xabcd);

#[derive(Clone)]
pub struct IPv4Packet {
    data: Vec<u8>,
}

impl IPv4Packet {
    pub fn construct<I1, I2>(source: I1, dest: I2, packet: &Layer4Packet) -> IPv4Packet
    where
        I1: Into<Ipv4Addr>,
        I2: Into<Ipv4Addr>,
    {
        let source = source.into();
        let dest = dest.into();

        let dscpen: u8 = 0;

        let totallen: u16 = 20 + packet.pkt().len();
        let id: u16 = IPV4_ID.fetch_add(1, Ordering::SeqCst);

        let fragmentation: u16 = 0x4000;

        let ttl: u8 = 64;
        let protocol: u8 = match packet {
            Layer4Packet::UDP(_) => 17,
        };

        let source_upper = u16::from_be_bytes(source.octets()[0..2].try_into().unwrap());
        let source_lower = u16::from_be_bytes(source.octets()[2..4].try_into().unwrap());
        let dest_upper = u16::from_be_bytes(dest.octets()[0..2].try_into().unwrap());
        let dest_lower = u16::from_be_bytes(dest.octets()[2..4].try_into().unwrap());

        let versihldscpen = 0x4500 | dscpen as u16;
        let ttlprotocol = ((ttl as u32) << 8_u8) | protocol as u32;

        let checksum_part: u32 = versihldscpen as u32
            + totallen as u32
            + id as u32
            + ttlprotocol
            + fragmentation as u32
            + source_upper as u32
            + source_lower as u32
            + dest_upper as u32
            + dest_lower as u32;

        let checksum: u16 = ((checksum_part & 0xFFFF) + (checksum_part >> 16))
            .try_into()
            .unwrap();

        let checksum = 0xFFFF - checksum;

        let data = [
            &versihldscpen.to_be_bytes()[..],
            &totallen.to_be_bytes(),
            &id.to_be_bytes(),
            &fragmentation.to_be_bytes(),
            &[ttl, protocol],
            &checksum.to_be_bytes(),
            &source.octets(),
            &dest.octets(),
            match packet {
                Layer4Packet::UDP(pkt) => &*pkt.data,
            },
        ]
        .concat();

        IPv4Packet { data }
    }

    #[inline]
    pub fn pkt(&'_ self) -> IPv4Pkt<'_> {
        IPv4Pkt { data: &self.data }
    }
}

#[derive(Clone)]
pub enum Layer4Packet {
    UDP(UDPPacket),
}

impl Layer4Packet {
    pub fn pkt(&'_ self) -> Layer4Pkt<'_> {
        match self {
            Layer4Packet::UDP(pkt) => Layer4Pkt::UDP(pkt.pkt()),
        }
    }
}

#[derive(Clone)]
pub struct UDPPacket {
    data: Vec<u8>,
}

impl UDPPacket {
    pub fn construct<T>(source: u16, dest: u16, data: T) -> UDPPacket
    where
        T: Into<Vec<u8>>,
    {
        let data = data.into();

        UDPPacket {
            data: [
                &source.to_be_bytes(),
                &dest.to_be_bytes(),
                &(8_u16 + TryInto::<u16>::try_into(data.len()).unwrap()).to_be_bytes(),
                &[0_u8, 0_u8],
                &*data,
            ]
            .concat(),
        }
    }

    #[inline]
    pub fn pkt(&'_ self) -> UDPPkt<'_> {
        UDPPkt { data: &self.data }
    }
}