use std::net::Ipv4Addr;

pub struct Ipv4Packet<'p> {
	/// The computed size of the IP packet
	header_length: usize,
	/// The total length of the IP packet when all fragments are combined
	total_length: u16,
	more_fragments: bool,
	fragment_offset: usize,
	next_header: IpNextHeader,
	pub src: Ipv4Addr,
	pub dst: Ipv4Addr,
	payload: &'p [u8],
}

impl<'p> Ipv4Packet<'p> {
	pub fn new(buffer: &'p [u8]) -> Self {
		let ihl = buffer[0].to_be() & 0b0000_1111;
		let header_length = ihl as usize * 4;

		let total_length = u16::from_be_bytes([buffer[2], buffer[3]]);

		// Fragmentation
		let more_fragments = (buffer[6] & 0b0010_0000) > 0;
		let fragment_offset = u16::from_be_bytes([buffer[6] & 0b0001_1111, buffer[7]]);
		// Fragments are in units of 8 bytes
		let true_frag_offset = fragment_offset as usize * 8;

		let next_header = IpNextHeader::new(buffer[9]);

		let src = Ipv4Addr::new(buffer[12], buffer[13], buffer[14], buffer[15]);
		let dst = Ipv4Addr::new(buffer[16], buffer[17], buffer[18], buffer[19]);

		Self {
			header_length,
			total_length,
			more_fragments,
			fragment_offset: true_frag_offset,
			next_header,
			src,
			dst,
			payload: &buffer[header_length..],
		}
	}

	pub fn get_source(&self) -> Ipv4Addr {
		self.src
	}

	pub fn get_destination(&self) -> Ipv4Addr {
		self.dst
	}

	pub fn has_more_fragments(&self) -> bool {
		self.more_fragments
	}

	pub fn get_fragment_offset(&self) -> usize {
		self.fragment_offset
	}

	pub fn get_next_header(&self) -> IpNextHeader {
		self.next_header
	}

	pub fn get_payload(&self) -> &[u8] {
		&self.payload
	}

	pub fn get_packet_len(&self) -> usize {
		self.total_length as usize
	}

	pub fn get_payload_len(&self) -> usize {
		// An IPv4 header is always 20 bytes long
		self.total_length as usize - 20
	}
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum IpNextHeader {
	Icmp,
	Igmp,
	Tcp,
	Udp,
	Unknown(u8),
}

impl IpNextHeader {
	pub fn new(n: u8) -> Self {
		match n {
			1 => Self::Icmp,
			2 => Self::Igmp,
			6 => Self::Tcp,
			17 => Self::Udp,
			n => Self::Unknown(n),
		}
	}
}