package bcap

import (
	"net"
	"testing"
	"time"

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

type observedPacket struct {
	Observation
	Key        string
	ReverseKey string
}

func wrapObservedPacket(obs Observation) observedPacket {
	return observedPacket{
		Observation: obs,
		Key:         obs.Flow.Forward.StableString(),
		ReverseKey:  obs.Flow.Reverse.StableString(),
	}
}

func (p observedPacket) StateDescript() uint8 {
	return legacyStateFromObservation(p.Observation)
}

func (p observedPacket) TcpPayloads() int {
	if tcp := p.Packet.Transport.TCP; tcp != nil {
		return tcp.Payload
	}
	return p.Packet.Transport.Payload
}

func (p observedPacket) TcpSeq() uint32 {
	if tcp := p.Packet.Transport.TCP; tcp != nil {
		return tcp.Seq
	}
	return 0
}

func (p observedPacket) TcpAck() uint32 {
	if tcp := p.Packet.Transport.TCP; tcp != nil {
		return tcp.Ack
	}
	return 0
}

func (p observedPacket) TcpWindow() uint16 {
	if tcp := p.Packet.Transport.TCP; tcp != nil {
		return tcp.Window
	}
	return 0
}

func legacyStateFromObservation(obs Observation) uint8 {
	switch obs.Packet.Transport.Kind {
	case ProtocolTCP:
		if obs.Hints.TCP != nil {
			return obs.Hints.TCP.LegacyState
		}
	case ProtocolUDP:
		return StateUdp
	case ProtocolICMPv4:
		return StateIcmp
	case ProtocolICMPv6:
		return StateIcmpv6
	}
	return StateUnknown
}

type tcpPacketSpec struct {
	srcIP   string
	dstIP   string
	srcPort uint16
	dstPort uint16
	seq     uint32
	ack     uint32
	syn     bool
	ackFlag bool
	fin     bool
	rst     bool
	ece     bool
	cwr     bool
	window  uint16
	payload []byte
}

type udpPacketSpec struct {
	srcIP   string
	dstIP   string
	srcPort uint16
	dstPort uint16
	payload []byte
}

type icmpPacketSpec struct {
	srcIP   string
	dstIP   string
	typ     uint8
	code    uint8
	id      uint16
	seq     uint16
	payload []byte
}

type arpPacketSpec struct {
	srcMAC    string
	dstMAC    string
	senderMAC string
	targetMAC string
	senderIP  string
	targetIP  string
	operation uint16
}

func newTestAnalyzer() *Analyzer {
	cfg := DefaultConfig()
	cfg.CleanupInterval = 0
	return NewAnalyzerWithConfig(cfg)
}

func mustObservePacket(t *testing.T, analyzer *Analyzer, packet gopacket.Packet) observedPacket {
	t.Helper()

	obs, err := analyzer.ObservePacket(packet)
	if err != nil {
		t.Fatalf("observe packet: %v", err)
	}
	return wrapObservedPacket(obs)
}

func mustEstablishTCPConnection(
	t *testing.T,
	analyzer *Analyzer,
	base time.Time,
	clientIP string,
	clientPort uint16,
	serverIP string,
	serverPort uint16,
	clientSeq uint32,
	serverSeq uint32,
) (string, string) {
	t.Helper()

	mustObservePacket(t, analyzer, mustBuildTCPPacket(t, base, tcpPacketSpec{
		srcIP:   clientIP,
		dstIP:   serverIP,
		srcPort: clientPort,
		dstPort: serverPort,
		seq:     clientSeq,
		syn:     true,
		window:  4096,
	}))

	serverInfo := mustObservePacket(t, analyzer, mustBuildTCPPacket(t, base.Add(time.Millisecond), tcpPacketSpec{
		srcIP:   serverIP,
		dstIP:   clientIP,
		srcPort: serverPort,
		dstPort: clientPort,
		seq:     serverSeq,
		ack:     clientSeq + 1,
		syn:     true,
		ackFlag: true,
		window:  4096,
	}))

	clientInfo := mustObservePacket(t, analyzer, mustBuildTCPPacket(t, base.Add(2*time.Millisecond), tcpPacketSpec{
		srcIP:   clientIP,
		dstIP:   serverIP,
		srcPort: clientPort,
		dstPort: serverPort,
		seq:     clientSeq + 1,
		ack:     serverSeq + 1,
		ackFlag: true,
		window:  4096,
	}))

	return clientInfo.Key, serverInfo.Key
}

func mustBuildTCPPacket(t *testing.T, ts time.Time, spec tcpPacketSpec) gopacket.Packet {
	t.Helper()

	ip := &layers.IPv4{
		Version:  4,
		IHL:      5,
		TTL:      64,
		SrcIP:    net.ParseIP(spec.srcIP).To4(),
		DstIP:    net.ParseIP(spec.dstIP).To4(),
		Protocol: layers.IPProtocolTCP,
	}
	tcp := &layers.TCP{
		SrcPort: layers.TCPPort(spec.srcPort),
		DstPort: layers.TCPPort(spec.dstPort),
		Seq:     spec.seq,
		Ack:     spec.ack,
		SYN:     spec.syn,
		ACK:     spec.ackFlag,
		FIN:     spec.fin,
		RST:     spec.rst,
		ECE:     spec.ece,
		CWR:     spec.cwr,
		Window:  spec.window,
	}
	tcp.SetNetworkLayerForChecksum(ip)

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err := gopacket.SerializeLayers(buf, opts, ip, tcp, gopacket.Payload(spec.payload)); err != nil {
		t.Fatalf("serialize tcp packet: %v", err)
	}

	packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
	packet.Metadata().Timestamp = ts
	return packet
}

func mustBuildUDPPacket(t *testing.T, ts time.Time, spec udpPacketSpec) gopacket.Packet {
	t.Helper()

	ip := &layers.IPv4{
		Version:  4,
		IHL:      5,
		TTL:      64,
		SrcIP:    net.ParseIP(spec.srcIP).To4(),
		DstIP:    net.ParseIP(spec.dstIP).To4(),
		Protocol: layers.IPProtocolUDP,
	}
	udp := &layers.UDP{
		SrcPort: layers.UDPPort(spec.srcPort),
		DstPort: layers.UDPPort(spec.dstPort),
	}
	udp.SetNetworkLayerForChecksum(ip)

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err := gopacket.SerializeLayers(buf, opts, ip, udp, gopacket.Payload(spec.payload)); err != nil {
		t.Fatalf("serialize udp packet: %v", err)
	}

	packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
	packet.Metadata().Timestamp = ts
	return packet
}

func mustBuildICMPPacket(t *testing.T, ts time.Time, spec icmpPacketSpec) gopacket.Packet {
	t.Helper()

	ip := &layers.IPv4{
		Version:  4,
		IHL:      5,
		TTL:      64,
		SrcIP:    net.ParseIP(spec.srcIP).To4(),
		DstIP:    net.ParseIP(spec.dstIP).To4(),
		Protocol: layers.IPProtocolICMPv4,
	}
	icmp := &layers.ICMPv4{
		TypeCode: layers.CreateICMPv4TypeCode(spec.typ, spec.code),
		Id:       spec.id,
		Seq:      spec.seq,
	}

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err := gopacket.SerializeLayers(buf, opts, ip, icmp, gopacket.Payload(spec.payload)); err != nil {
		t.Fatalf("serialize icmp packet: %v", err)
	}

	packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
	packet.Metadata().Timestamp = ts
	return packet
}

func mustBuildIPv4OnlyPacket(t *testing.T, ts time.Time, srcIP, dstIP string) gopacket.Packet {
	t.Helper()

	ip := &layers.IPv4{
		Version:  4,
		IHL:      5,
		TTL:      64,
		SrcIP:    net.ParseIP(srcIP).To4(),
		DstIP:    net.ParseIP(dstIP).To4(),
		Protocol: layers.IPProtocolNoNextHeader,
	}

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err := gopacket.SerializeLayers(buf, opts, ip); err != nil {
		t.Fatalf("serialize ipv4-only packet: %v", err)
	}

	packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeIPv4, gopacket.Default)
	packet.Metadata().Timestamp = ts
	return packet
}

func mustBuildARPPacket(t *testing.T, ts time.Time, spec arpPacketSpec) gopacket.Packet {
	t.Helper()

	eth := &layers.Ethernet{
		SrcMAC:       mustMAC(t, spec.srcMAC),
		DstMAC:       mustMAC(t, spec.dstMAC),
		EthernetType: layers.EthernetTypeARP,
	}
	arp := &layers.ARP{
		AddrType:          layers.LinkTypeEthernet,
		Protocol:          layers.EthernetTypeIPv4,
		HwAddressSize:     6,
		ProtAddressSize:   4,
		Operation:         spec.operation,
		SourceHwAddress:   []byte(mustMAC(t, spec.senderMAC)),
		SourceProtAddress: []byte(net.ParseIP(spec.senderIP).To4()),
		DstHwAddress:      []byte(mustMAC(t, spec.targetMAC)),
		DstProtAddress:    []byte(net.ParseIP(spec.targetIP).To4()),
	}

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, arp); err != nil {
		t.Fatalf("serialize arp packet: %v", err)
	}

	packet := gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
	packet.Metadata().Timestamp = ts
	return packet
}

func mustMAC(t *testing.T, raw string) net.HardwareAddr {
	t.Helper()

	hw, err := net.ParseMAC(raw)
	if err != nil {
		t.Fatalf("parse mac %q: %v", raw, err)
	}
	return hw
}