package aerotun

import (
	"bytes"
	"encoding/binary"
	"testing"
)

func buildMockIPv4Packet(proto byte, srcPort, dstPort uint16, payloadLen int) []byte {
	hdrLen := 20
	pkt := make([]byte, hdrLen+4+payloadLen)
	pkt[0] = 0x45 // IPv4, IHL=5 (20 bytes)
	binary.BigEndian.PutUint16(pkt[2:4], uint16(len(pkt)))
	pkt[9] = proto
	copy(pkt[12:16], []byte{10, 99, 0, 1})
	copy(pkt[16:20], []byte{10, 99, 0, 2})
	if proto == 6 || proto == 17 {
		binary.BigEndian.PutUint16(pkt[20:22], srcPort)
		binary.BigEndian.PutUint16(pkt[22:24], dstPort)
	}
	return pkt
}

func buildMockIPv6Packet(proto byte, srcPort, dstPort uint16, payloadLen int) []byte {
	pkt := make([]byte, 40+4+payloadLen)
	pkt[0] = 0x60 // IPv6
	binary.BigEndian.PutUint16(pkt[4:6], uint16(4+payloadLen))
	pkt[6] = proto
	if proto == 6 || proto == 17 {
		binary.BigEndian.PutUint16(pkt[40:42], srcPort)
		binary.BigEndian.PutUint16(pkt[42:44], dstPort)
	}
	return pkt
}

func TestIsEncryptedIPPacket(t *testing.T) {
	// IPv4 HTTPS (port 443)
	pkt4HTTPS := buildMockIPv4Packet(6, 50123, 443, 200)
	if !isEncryptedIPPacket(pkt4HTTPS) {
		t.Fatalf("expected IPv4 port 443 to be recognized as encrypted")
	}

	// IPv4 WireGuard (port 51820)
	pkt4WG := buildMockIPv4Packet(17, 51820, 50124, 200)
	if !isEncryptedIPPacket(pkt4WG) {
		t.Fatalf("expected IPv4 port 51820 to be recognized as encrypted")
	}

	// IPv4 HTTP (port 80) - Plaintext
	pkt4HTTP := buildMockIPv4Packet(6, 50125, 80, 200)
	if isEncryptedIPPacket(pkt4HTTP) {
		t.Fatalf("expected IPv4 port 80 to not be recognized as encrypted")
	}

	// IPv6 SSH (port 22)
	pkt6SSH := buildMockIPv6Packet(6, 22, 50126, 200)
	if !isEncryptedIPPacket(pkt6SSH) {
		t.Fatalf("expected IPv6 port 22 to be recognized as encrypted")
	}

	// IPv6 DNS (port 53) - Plaintext
	pkt6DNS := buildMockIPv6Packet(17, 50127, 53, 200)
	if isEncryptedIPPacket(pkt6DNS) {
		t.Fatalf("expected IPv6 port 53 to not be recognized as encrypted")
	}
}

func TestAdaptiveCompressor(t *testing.T) {
	ac := newAdaptiveCompressor(1420)

	// 1. Small packet (< 128 bytes) should be bypassed
	smallPkt := buildMockIPv4Packet(6, 50128, 80, 50)
	if _, ok := ac.Compress(smallPkt); ok {
		t.Fatalf("small packet should not be compressed")
	}

	// 2. Encrypted packet (port 443) should be bypassed without LZ4
	encPkt := buildMockIPv4Packet(6, 50129, 443, 300)
	if _, ok := ac.Compress(encPkt); ok {
		t.Fatalf("encrypted packet should not be compressed")
	}

	// 3. Compressible packet (e.g. repeated pattern) on port 80
	compressiblePayload := bytes.Repeat([]byte("ABCDEFGHIJ1234567890"), 30) // 600 bytes
	compPkt := buildMockIPv4Packet(6, 50130, 80, len(compressiblePayload))
	copy(compPkt[24:], compressiblePayload)

	out, ok := ac.Compress(compPkt)
	if !ok || len(out) >= len(compPkt)-16 {
		t.Fatalf("expected compressible packet to compress, got ok=%v, len=%d", ok, len(out))
	}

	// 4. Incompressible random packet stream triggers backoff
	incompPayload := []byte{
		52, 56, 184, 144, 222, 129, 243, 220, 31, 151, 230, 149, 61, 206, 167, 169, 127, 181, 162, 174,
		132, 186, 105, 47, 180, 106, 222, 232, 2, 220, 117, 243, 3, 5, 26, 229, 136, 49, 9, 249,
		92, 89, 189, 80, 114, 199, 93, 244, 95, 21, 15, 212, 5, 241, 24, 101, 13, 94, 151, 67,
		88, 236, 222, 216, 162, 133, 240, 45, 138, 47, 11, 225, 207, 45, 7, 99, 13, 23, 213, 53,
		212, 12, 113, 51, 97, 247, 58, 49, 201, 108, 243, 54, 238, 153, 40, 189, 66, 149, 24, 83,
		8, 184, 76, 123, 233, 190, 68, 207, 249, 87, 16, 54, 121, 80, 188, 45, 198, 21, 254, 107,
		159, 52, 187, 220, 11, 159, 180, 235, 219, 209, 192, 60, 101, 201, 195, 137, 109, 195, 70, 124,
		65, 79, 209, 48, 184, 253, 16, 132, 118, 15, 104, 50, 114, 157, 121, 62, 155, 141, 77, 82,
		222, 58, 128, 0, 84, 157, 91, 156, 189, 239, 242, 109, 78, 187, 221, 204, 35, 201, 158, 94,
		20, 199, 103, 136, 115, 17, 114, 255, 115, 46, 182, 194, 41, 25, 104, 56, 154, 244, 134, 158,
	}
	incompPkt := buildMockIPv4Packet(6, 50131, 80, len(incompPayload))
	copy(incompPkt[24:], incompPayload)

	// Fire 16 failures
	for i := 0; i < 16; i++ {
		_, ok := ac.Compress(incompPkt)
		if ok {
			t.Fatalf("random packet should not compress")
		}
	}

	// At this point, backoff should be triggered (skipCount == 256)
	if ac.skipCount != 256 {
		t.Fatalf("expected skipCount == 256, got %d", ac.skipCount)
	}

	// The next 256 calls should be skipped immediately
	for i := 0; i < 256; i++ {
		_, ok := ac.Compress(incompPkt)
		if ok {
			t.Fatalf("incompressible packet should be skipped during backoff")
		}
	}
	if ac.skipCount != 0 {
		t.Fatalf("expected skipCount to count down to 0, got %d", ac.skipCount)
	}
}