#include <criterion/criterion.h>
#include <string.h>
#include <arpa/inet.h>
#include "internal/icmp_packet_internal.h"
#include "internal/icmp_packet.h"

Test(packet_parse_ip, valid_header)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45; // version 4, IHL = 5*4 = 20 bytes
	hdr->ttl = 64;
	hdr->saddr = inet_addr("192.168.1.1");

	uint8_t ttl = 0;
	struct in_addr src_addr = {0};
	size_t ip_hdr_len = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), &ttl, &src_addr,
									&ip_hdr_len, NULL, NULL);

	cr_assert_eq(ret, 0, "Should return 0 for valid IP header");
	cr_assert_eq(ttl, 64, "TTL should be 64");
	cr_assert_eq(src_addr.s_addr, inet_addr("192.168.1.1"),
				"Source address mismatch");
	cr_assert_eq(ip_hdr_len, 20, "Header length should be 20 bytes");
}

Test(packet_parse_ip, extract_ttl)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45;
	hdr->ttl = 128;
	hdr->saddr = inet_addr("10.0.0.1");

	uint8_t ttl = 0;
	struct in_addr src_addr = {0};
	size_t ip_hdr_len = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), &ttl, &src_addr,
									&ip_hdr_len, NULL, NULL);

	cr_assert_eq(ret, 0, "Should return 0 for valid header");
	cr_assert_eq(ttl, 128, "TTL should be 128");
}

Test(packet_parse_ip, extract_source_addr)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45;
	hdr->ttl = 5;
	hdr->saddr = inet_addr("8.8.8.8");

	uint8_t ttl = 0;
	struct in_addr src_addr = {0};
	size_t ip_hdr_len = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), &ttl, &src_addr,
									&ip_hdr_len, NULL, NULL);

	cr_assert_eq(ret, 0, "Should return 0 for valid header");
	cr_assert_eq(src_addr.s_addr, inet_addr("8.8.8.8"),
				"Source address should match 8.8.8.8");
}

Test(packet_parse_ip, buffer_too_small)
{
	uint8_t buffer[10] = {0}; // < 20
	uint8_t ttl = 0;
	struct in_addr src_addr = {0};
	size_t ip_hdr_len = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), &ttl, &src_addr,
									&ip_hdr_len, NULL, NULL);
	cr_assert_eq(ret, -1, "Should return -1 when buffer is too small");

}

Test(packet_parse_ip, invalid_version)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x65; // version = 6, IHL = 5
	hdr->ttl = 99;
	hdr->saddr = inet_addr("1.2.3.4");

	uint8_t ttl = 0;
	struct in_addr src_addr = {0};
	size_t ip_hdr_len = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), &ttl, &src_addr,
									&ip_hdr_len, NULL, NULL);
	cr_assert_eq(ret, -1, "Should return -1 for invalid IP version");
}

Test(packet_parse_ip, extract_dest_addr)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45;
	hdr->ttl = 64;
	hdr->saddr = inet_addr("192.168.1.1");
	hdr->daddr = inet_addr("8.8.8.8");

	struct in_addr dst_addr = {0};

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), NULL, NULL,
									NULL, &dst_addr, NULL);

	cr_assert_eq(ret, 0, "Should return 0 for valid header");
	cr_assert_eq(dst_addr.s_addr, inet_addr("8.8.8.8"),
				"Destination address should match 8.8.8.8");
}

Test(packet_parse_ip, extract_protocol)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45;
	hdr->ttl = 64;
	hdr->protocol = 1; // ICMP

	uint8_t protocol = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), NULL, NULL,
									NULL, NULL, &protocol);

	cr_assert_eq(ret, 0, "Should return 0 for valid header");
	cr_assert_eq(protocol, 1, "Protocol should be 1 (ICMP)");
}

Test(packet_parse_ip, extract_all_fields)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45;
	hdr->ttl = 128;
	hdr->saddr = inet_addr("10.0.0.1");
	hdr->daddr = inet_addr("10.0.0.2");
	hdr->protocol = 6; // TCP

	uint8_t ttl = 0;
	struct in_addr src_addr = {0};
	struct in_addr dst_addr = {0};
	uint8_t protocol = 0;
	size_t ip_hdr_len = 0;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), &ttl, &src_addr,
									&ip_hdr_len, &dst_addr, &protocol);

	cr_assert_eq(ret, 0, "Should return 0 for valid header");
	cr_assert_eq(ttl, 128, "TTL should be 128");
	cr_assert_eq(src_addr.s_addr, inet_addr("10.0.0.1"),
				"Source address should match");
	cr_assert_eq(dst_addr.s_addr, inet_addr("10.0.0.2"),
				"Destination address should match");
	cr_assert_eq(protocol, 6, "Protocol should be 6 (TCP)");
	cr_assert_eq(ip_hdr_len, 20, "Header length should be 20");
}

Test(packet_parse_ip, null_parameters_allowed)
{
	uint8_t buffer[20] = {0};
	struct ip_header *hdr = (struct ip_header*)buffer;
	hdr->version_ihl = 0x45;
	hdr->ttl = 64;

	int ret = icmp_parse_ip_header(buffer, sizeof(buffer), NULL, NULL,
									NULL, NULL, NULL);

	cr_assert_eq(ret, 0, "Should return 0 even with all NULL parameters");
}