package classify_test

import (
	"strings"
	"testing"

	"github.com/gopacket/gopacket"
	"github.com/gopacket/gopacket/layers"
	"github.com/netsynth/netsynth/classify"
)

// buildTCPPacket builds a synthetic IPv4/TCP packet targeting the given dst port.
func buildTCPPacket(t *testing.T, dstPort uint16) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: false}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeIPv4,
	}
	ip := &layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolTCP,
		SrcIP:    []byte{192, 168, 1, 1},
		DstIP:    []byte{192, 168, 1, 2},
	}
	tcp := &layers.TCP{
		SrcPort: layers.TCPPort(54321),
		DstPort: layers.TCPPort(dstPort),
		SYN:     true,
	}
	if err := tcp.SetNetworkLayerForChecksum(ip); err != nil {
		t.Fatalf("SetNetworkLayerForChecksum: %v", err)
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, ip, tcp, gopacket.Payload{}); err != nil {
		t.Fatalf("SerializeLayers tcp: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

// buildUDPPacket builds a synthetic IPv4/UDP packet targeting the given dst port.
func buildUDPPacket(t *testing.T, dstPort uint16) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: false}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeIPv4,
	}
	ip := &layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolUDP,
		SrcIP:    []byte{192, 168, 1, 1},
		DstIP:    []byte{192, 168, 1, 2},
	}
	udp := &layers.UDP{
		SrcPort: layers.UDPPort(54321),
		DstPort: layers.UDPPort(dstPort),
	}
	if err := udp.SetNetworkLayerForChecksum(ip); err != nil {
		t.Fatalf("SetNetworkLayerForChecksum: %v", err)
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, ip, udp, gopacket.Payload{}); err != nil {
		t.Fatalf("SerializeLayers udp: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

// buildICMPPacket builds a synthetic IPv4/ICMPv4 packet.
func buildICMPPacket(t *testing.T) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: false}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeIPv4,
	}
	ip := &layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolICMPv4,
		SrcIP:    []byte{192, 168, 1, 1},
		DstIP:    []byte{192, 168, 1, 2},
	}
	icmp := &layers.ICMPv4{
		TypeCode: layers.CreateICMPv4TypeCode(layers.ICMPv4TypeEchoRequest, 0),
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, ip, icmp, gopacket.Payload{}); err != nil {
		t.Fatalf("SerializeLayers icmp: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

// buildUnknownPacket builds a synthetic Ethernet packet with no TCP/UDP/ICMP payload.
func buildUnknownPacket(t *testing.T) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeARP,
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, gopacket.Payload{0x01, 0x02, 0x03, 0x04}); err != nil {
		t.Fatalf("SerializeLayers unknown: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

func TestClassify(t *testing.T) {
	c := classify.NewClassifier(classify.DefaultRules)

	t.Run("TestClassifyICMP", func(t *testing.T) {
		pkt := buildICMPPacket(t)
		got := c.Classify(pkt)
		if got.Class != classify.ClassICMP {
			t.Errorf("ICMP packet: got class %q, want %q", got.Class, classify.ClassICMP)
		}
		if got.Protocol != "icmp" {
			t.Errorf("ICMP packet: got protocol %q, want %q", got.Protocol, "icmp")
		}
	})

	t.Run("TestClassifyDNS_UDP", func(t *testing.T) {
		pkt := buildUDPPacket(t, 53)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDNS {
			t.Errorf("DNS/UDP packet: got class %q, want %q", got.Class, classify.ClassDNS)
		}
	})

	t.Run("TestClassifyDNS_TCP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 53)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDNS {
			t.Errorf("DNS/TCP packet: got class %q, want %q", got.Class, classify.ClassDNS)
		}
	})

	t.Run("TestClassifyHTTPS", func(t *testing.T) {
		pkt := buildTCPPacket(t, 443)
		got := c.Classify(pkt)
		if got.Class != classify.ClassHTTPS {
			t.Errorf("HTTPS packet: got class %q, want %q", got.Class, classify.ClassHTTPS)
		}
	})

	t.Run("TestClassifyHTTP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 80)
		got := c.Classify(pkt)
		if got.Class != classify.ClassHTTP {
			t.Errorf("HTTP packet: got class %q, want %q", got.Class, classify.ClassHTTP)
		}
	})

	t.Run("TestClassifySSH", func(t *testing.T) {
		pkt := buildTCPPacket(t, 22)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSSH {
			t.Errorf("SSH packet: got class %q, want %q", got.Class, classify.ClassSSH)
		}
	})

	t.Run("TestClassifySMTP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 25)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSMTP {
			t.Errorf("SMTP packet: got class %q, want %q", got.Class, classify.ClassSMTP)
		}
	})

	t.Run("TestClassifyNTP", func(t *testing.T) {
		pkt := buildUDPPacket(t, 123)
		got := c.Classify(pkt)
		if got.Class != classify.ClassNTP {
			t.Errorf("NTP packet: got class %q, want %q", got.Class, classify.ClassNTP)
		}
	})

	t.Run("TestClassifyDHCP_port67", func(t *testing.T) {
		pkt := buildUDPPacket(t, 67)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDHCP {
			t.Errorf("DHCP port 67 packet: got class %q, want %q", got.Class, classify.ClassDHCP)
		}
	})

	t.Run("TestClassifyDHCP_port68", func(t *testing.T) {
		pkt := buildUDPPacket(t, 68)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDHCP {
			t.Errorf("DHCP port 68 packet: got class %q, want %q", got.Class, classify.ClassDHCP)
		}
	})

	t.Run("TestClassifyOtherTCP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 8080)
		got := c.Classify(pkt)
		if got.Class != classify.ClassOtherTCP {
			t.Errorf("OtherTCP packet (port 8080): got class %q, want %q", got.Class, classify.ClassOtherTCP)
		}
	})

	t.Run("TestClassifyOtherUDP", func(t *testing.T) {
		pkt := buildUDPPacket(t, 9999)
		got := c.Classify(pkt)
		if got.Class != classify.ClassOtherUDP {
			t.Errorf("OtherUDP packet (port 9999): got class %q, want %q", got.Class, classify.ClassOtherUDP)
		}
	})

	t.Run("TestClassifyUnknown", func(t *testing.T) {
		pkt := buildUnknownPacket(t)
		got := c.Classify(pkt)
		if !strings.HasPrefix(string(got.Class), "unknown-") {
			t.Errorf("Unknown packet: got class %q, want unknown-N bucket", got.Class)
		}
	})

	// --- Mail (PROTO-01) ---
	t.Run("TestClassifyIMAP_port143", func(t *testing.T) {
		pkt := buildTCPPacket(t, 143)
		got := c.Classify(pkt)
		if got.Class != classify.ClassIMAP {
			t.Errorf("IMAP port 143: got class %q, want %q", got.Class, classify.ClassIMAP)
		}
	})
	t.Run("TestClassifyIMAP_port993", func(t *testing.T) {
		pkt := buildTCPPacket(t, 993)
		got := c.Classify(pkt)
		if got.Class != classify.ClassIMAP {
			t.Errorf("IMAPS port 993: got class %q, want %q", got.Class, classify.ClassIMAP)
		}
	})
	t.Run("TestClassifyPOP3_port110", func(t *testing.T) {
		pkt := buildTCPPacket(t, 110)
		got := c.Classify(pkt)
		if got.Class != classify.ClassPOP3 {
			t.Errorf("POP3 port 110: got class %q, want %q", got.Class, classify.ClassPOP3)
		}
	})
	t.Run("TestClassifyPOP3_port995", func(t *testing.T) {
		pkt := buildTCPPacket(t, 995)
		got := c.Classify(pkt)
		if got.Class != classify.ClassPOP3 {
			t.Errorf("POP3S port 995: got class %q, want %q", got.Class, classify.ClassPOP3)
		}
	})
	t.Run("TestClassifySMTPSub_port587", func(t *testing.T) {
		pkt := buildTCPPacket(t, 587)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSMTPSub {
			t.Errorf("SMTP-sub port 587: got class %q, want %q", got.Class, classify.ClassSMTPSub)
		}
	})

	// --- File Transfer (PROTO-02) ---
	t.Run("TestClassifyFTP_port20", func(t *testing.T) {
		pkt := buildTCPPacket(t, 20)
		got := c.Classify(pkt)
		if got.Class != classify.ClassFTP {
			t.Errorf("FTP data port 20: got class %q, want %q", got.Class, classify.ClassFTP)
		}
	})
	t.Run("TestClassifyFTP_port21", func(t *testing.T) {
		pkt := buildTCPPacket(t, 21)
		got := c.Classify(pkt)
		if got.Class != classify.ClassFTP {
			t.Errorf("FTP control port 21: got class %q, want %q", got.Class, classify.ClassFTP)
		}
	})
	t.Run("TestClassifySMB_port445", func(t *testing.T) {
		pkt := buildTCPPacket(t, 445)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSMB {
			t.Errorf("SMB port 445: got class %q, want %q", got.Class, classify.ClassSMB)
		}
	})
	t.Run("TestClassifyTFTP_port69", func(t *testing.T) {
		pkt := buildUDPPacket(t, 69)
		got := c.Classify(pkt)
		if got.Class != classify.ClassTFTP {
			t.Errorf("TFTP port 69: got class %q, want %q", got.Class, classify.ClassTFTP)
		}
	})

	// --- Remote Access (PROTO-03) ---
	t.Run("TestClassifyRDP_port3389", func(t *testing.T) {
		pkt := buildTCPPacket(t, 3389)
		got := c.Classify(pkt)
		if got.Class != classify.ClassRDP {
			t.Errorf("RDP port 3389: got class %q, want %q", got.Class, classify.ClassRDP)
		}
	})
	t.Run("TestClassifyTelnet_port23", func(t *testing.T) {
		pkt := buildTCPPacket(t, 23)
		got := c.Classify(pkt)
		if got.Class != classify.ClassTelnet {
			t.Errorf("Telnet port 23: got class %q, want %q", got.Class, classify.ClassTelnet)
		}
	})
	t.Run("TestClassifyVNC_port5900", func(t *testing.T) {
		pkt := buildTCPPacket(t, 5900)
		got := c.Classify(pkt)
		if got.Class != classify.ClassVNC {
			t.Errorf("VNC port 5900: got class %q, want %q", got.Class, classify.ClassVNC)
		}
	})

	// --- Database (PROTO-04) ---
	t.Run("TestClassifyMySQL_port3306", func(t *testing.T) {
		pkt := buildTCPPacket(t, 3306)
		got := c.Classify(pkt)
		if got.Class != classify.ClassMySQL {
			t.Errorf("MySQL port 3306: got class %q, want %q", got.Class, classify.ClassMySQL)
		}
	})
	t.Run("TestClassifyPostgreSQL_port5432", func(t *testing.T) {
		pkt := buildTCPPacket(t, 5432)
		got := c.Classify(pkt)
		if got.Class != classify.ClassPostgreSQL {
			t.Errorf("PostgreSQL port 5432: got class %q, want %q", got.Class, classify.ClassPostgreSQL)
		}
	})
	t.Run("TestClassifyRedis_port6379", func(t *testing.T) {
		pkt := buildTCPPacket(t, 6379)
		got := c.Classify(pkt)
		if got.Class != classify.ClassRedis {
			t.Errorf("Redis port 6379: got class %q, want %q", got.Class, classify.ClassRedis)
		}
	})
	t.Run("TestClassifyMongoDB_port27017", func(t *testing.T) {
		pkt := buildTCPPacket(t, 27017)
		got := c.Classify(pkt)
		if got.Class != classify.ClassMongoDB {
			t.Errorf("MongoDB port 27017: got class %q, want %q", got.Class, classify.ClassMongoDB)
		}
	})

	// --- Discovery (PROTO-05) ---
	t.Run("TestClassifyMDNS_port5353", func(t *testing.T) {
		pkt := buildUDPPacket(t, 5353)
		got := c.Classify(pkt)
		if got.Class != classify.ClassMDNS {
			t.Errorf("mDNS port 5353: got class %q, want %q", got.Class, classify.ClassMDNS)
		}
	})
	t.Run("TestClassifySDP_port1900", func(t *testing.T) {
		pkt := buildUDPPacket(t, 1900)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSSDP {
			t.Errorf("SSDP port 1900: got class %q, want %q", got.Class, classify.ClassSSDP)
		}
	})
	t.Run("TestClassifySNMP_port161", func(t *testing.T) {
		pkt := buildUDPPacket(t, 161)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSNMP {
			t.Errorf("SNMP port 161: got class %q, want %q", got.Class, classify.ClassSNMP)
		}
	})
	t.Run("TestClassifySNMP_port162", func(t *testing.T) {
		pkt := buildUDPPacket(t, 162)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSNMP {
			t.Errorf("SNMP-trap port 162: got class %q, want %q", got.Class, classify.ClassSNMP)
		}
	})

	// --- VoIP (PROTO-06) --- SIP runs on both TCP and UDP
	t.Run("TestClassifySIP_TCP5060", func(t *testing.T) {
		pkt := buildTCPPacket(t, 5060)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSIP {
			t.Errorf("SIP TCP 5060: got class %q, want %q", got.Class, classify.ClassSIP)
		}
	})
	t.Run("TestClassifySIP_UDP5060", func(t *testing.T) {
		pkt := buildUDPPacket(t, 5060)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSIP {
			t.Errorf("SIP UDP 5060: got class %q, want %q", got.Class, classify.ClassSIP)
		}
	})

	// --- Web extension (PROTO-07) --- QUIC is UDP 443, must NOT match HTTPS (TCP 443)
	t.Run("TestClassifyQUIC_UDP443", func(t *testing.T) {
		pkt := buildUDPPacket(t, 443)
		got := c.Classify(pkt)
		if got.Class != classify.ClassQUIC {
			t.Errorf("QUIC UDP 443: got class %q, want %q", got.Class, classify.ClassQUIC)
		}
	})

	// --- Infrastructure extension (PROTO-08) ---
	t.Run("TestClassifyLDAP_port389", func(t *testing.T) {
		pkt := buildTCPPacket(t, 389)
		got := c.Classify(pkt)
		if got.Class != classify.ClassLDAP {
			t.Errorf("LDAP port 389: got class %q, want %q", got.Class, classify.ClassLDAP)
		}
	})
	t.Run("TestClassifyLDAP_port636", func(t *testing.T) {
		pkt := buildTCPPacket(t, 636)
		got := c.Classify(pkt)
		if got.Class != classify.ClassLDAP {
			t.Errorf("LDAPS port 636: got class %q, want %q", got.Class, classify.ClassLDAP)
		}
	})
	t.Run("TestClassifyKerberos_TCP88", func(t *testing.T) {
		pkt := buildTCPPacket(t, 88)
		got := c.Classify(pkt)
		if got.Class != classify.ClassKerberos {
			t.Errorf("Kerberos TCP 88: got class %q, want %q", got.Class, classify.ClassKerberos)
		}
	})
	t.Run("TestClassifyKerberos_UDP88", func(t *testing.T) {
		pkt := buildUDPPacket(t, 88)
		got := c.Classify(pkt)
		if got.Class != classify.ClassKerberos {
			t.Errorf("Kerberos UDP 88: got class %q, want %q", got.Class, classify.ClassKerberos)
		}
	})
	t.Run("TestClassifySyslog_port514", func(t *testing.T) {
		pkt := buildUDPPacket(t, 514)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSyslog {
			t.Errorf("Syslog port 514: got class %q, want %q", got.Class, classify.ClassSyslog)
		}
	})

	t.Run("TestRulesAreOrderDependent", func(t *testing.T) {
		// A rule list where TCP port 443 maps to ClassHTTP (wrong) placed first,
		// then ClassHTTPS. The first matching rule should win.
		customRules := []classify.Rule{
			{Protocol: "tcp", DstPort: 443, Class: classify.ClassHTTP}, // first match
			{Protocol: "tcp", DstPort: 443, Class: classify.ClassHTTPS},
			{Protocol: "tcp", DstPort: 0, Class: classify.ClassOtherTCP},
		}
		custom := classify.NewClassifier(customRules)
		pkt := buildTCPPacket(t, 443)
		got := custom.Classify(pkt)
		if got.Class != classify.ClassHTTP {
			t.Errorf("Order-dependent rules: got class %q, want %q (first match should win)", got.Class, classify.ClassHTTP)
		}
	})
}

func TestAllClassesCount(t *testing.T) {
	classes := classify.AllClasses()
	if len(classes) != 32 {
		t.Errorf("AllClasses() returned %d classes, want 32", len(classes))
	}
	// No plain "unknown" should exist
	for _, c := range classes {
		if c == "unknown" {
			t.Error("AllClasses() still contains plain \"unknown\" — should be removed per D-04")
		}
	}
	// All 4 buckets must be present
	buckets := map[classify.TrafficClass]bool{
		classify.ClassUnknown1: false, classify.ClassUnknown2: false,
		classify.ClassUnknown3: false, classify.ClassUnknown4: false,
	}
	for _, c := range classes {
		if _, ok := buckets[c]; ok {
			buckets[c] = true
		}
	}
	for bucket, found := range buckets {
		if !found {
			t.Errorf("AllClasses() missing bucket %q", bucket)
		}
	}
}

func TestHashBucketDistribution(t *testing.T) {
	// hashBucket is unexported, test through Classify with a minimal rule set
	// that has no catch-all TCP/UDP rules, so unmatched ports reach hashBucket.
	minimalRules := []classify.Rule{
		{Protocol: "icmp", DstPort: 0, Class: classify.ClassICMP},
		{Protocol: "tcp", DstPort: 443, Class: classify.ClassHTTPS},
	}
	c := classify.NewClassifier(minimalRules)
	seen := make(map[classify.TrafficClass]bool)
	// Try a range of unmatched TCP ports to hit all 4 buckets
	for port := uint16(10000); port < 11000; port++ {
		pkt := buildTCPPacket(t, port)
		got := c.Classify(pkt)
		if strings.HasPrefix(string(got.Class), "unknown-") {
			seen[got.Class] = true
		}
	}
	for _, bucket := range []classify.TrafficClass{
		classify.ClassUnknown1, classify.ClassUnknown2,
		classify.ClassUnknown3, classify.ClassUnknown4,
	} {
		if !seen[bucket] {
			t.Errorf("hashBucket never produced %q across ports 10000-10999", bucket)
		}
	}
}