package main

import (
	"net/netip"
	"testing"
	"time"

	"github.com/ti-mo/conntrack"
)

func TestConntrackRecordRequiresDeclaredPortAndReplyTarget(t *testing.T) {
	flow := acceptedFlow()
	record, ok := conntrackRecordForFlow(&flow, []conntrackService{{
		Name:          "web",
		Ports:         []uint16{80, 443},
		TargetAddress: netip.MustParseAddr("10.104.40.127"),
	}})
	if !ok {
		t.Fatal("accepted relay flow was not classified")
	}
	if record.Service != "web" || record.Original.SourceAddress != "198.18.0.1" || record.Reply.SourceAddress != "10.104.40.127" || record.DurationMilliseconds != 1500 || record.TCPState != "time_wait" {
		t.Fatalf("unexpected record %#v", record)
	}
	flow.TupleReply.IP.SourceAddress = netip.MustParseAddr("203.0.113.10")
	if _, ok := conntrackRecordForFlow(&flow, []conntrackService{{Name: "web", Ports: []uint16{443}, TargetAddress: netip.MustParseAddr("10.104.40.127")}}); ok {
		t.Fatal("outbound connection was classified as accepted relay traffic")
	}
}

func TestConntrackSamplingIsDeterministic(t *testing.T) {
	flow := acceptedFlow()
	first := sampleConntrackFlow(&flow, 10)
	for attempt := 0; attempt < 100; attempt++ {
		if sampleConntrackFlow(&flow, 10) != first {
			t.Fatal("sampling decision changed for an unchanged tuple")
		}
	}
}

func TestTokenBucketBoundsBurstAndRefills(t *testing.T) {
	now := time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC)
	bucket := newTokenBucket(1, 2, now)
	if !bucket.allow(now) || !bucket.allow(now) || bucket.allow(now) {
		t.Fatal("initial burst was not bounded")
	}
	if !bucket.allow(now.Add(time.Second)) || bucket.allow(now.Add(time.Second)) {
		t.Fatal("sustained rate was not bounded")
	}
}

func acceptedFlow() conntrack.Flow {
	startedAt := time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC)
	flow := conntrack.NewFlow(6, 0, netip.MustParseAddr("198.18.0.1"), netip.MustParseAddr("10.49.127.43"), 50000, 443, 60, 0)
	flow.TupleReply.IP.SourceAddress = netip.MustParseAddr("10.104.40.127")
	flow.TupleReply.IP.DestinationAddress = netip.MustParseAddr("198.18.0.1")
	flow.TupleReply.Proto.SourcePort = 443
	flow.TupleReply.Proto.DestinationPort = 50000
	flow.Timestamp = conntrack.Timestamp{Start: startedAt, Stop: startedAt.Add(1500 * time.Millisecond)}
	flow.CountersOrig = conntrack.Counter{Packets: 3, Bytes: 180}
	flow.CountersReply = conntrack.Counter{Direction: true, Packets: 4, Bytes: 240}
	flow.ProtoInfo.TCP = &conntrack.ProtoInfoTCP{State: 7}
	return flow
}