package pcap

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"math/rand/v2"
	"slices"
	"strings"
	"testing"
	"time"
)

// reorderTestFile builds a pcap file with one packet per timestamp; each
// packet's data is its index in the input.
func reorderTestFile(t *testing.T, ts ...time.Time) []byte {
	t.Helper()
	var buf bytes.Buffer
	w, err := NewWriter(&buf, LinkTypeEthernet)
	if err != nil {
		t.Fatal(err)
	}
	for i, tm := range ts {
		if err := w.WritePacket(tm, bytes.Repeat([]byte{byte(i)}, i+1)); err != nil {
			t.Fatal(err)
		}
	}
	return buf.Bytes()
}

// packetIDs returns the first data byte of each packet in a pcap file.
func packetIDs(t *testing.T, file []byte) []int {
	t.Helper()
	r, err := NewReader(bytes.NewReader(file))
	if err != nil {
		t.Fatal(err)
	}
	var ids []int
	for {
		rec, err := r.Next()
		if err == io.EOF {
			return ids
		}
		if err != nil {
			t.Fatal(err)
		}
		ids = append(ids, int(rec[16]))
	}
}

func TestReorder(t *testing.T) {
	s := func(sec int64, usec int) time.Time { return time.Unix(sec, int64(usec)*1000) }
	tests := []struct {
		name       string
		ts         []time.Time
		want       []int
		outOfOrder int
	}{
		{"sorted", []time.Time{s(1, 0), s(1, 1), s(2, 0)}, []int{0, 1, 2}, 0},
		{"reversed", []time.Time{s(3, 0), s(2, 0), s(1, 0)}, []int{2, 1, 0}, 2},
		{"microseconds", []time.Time{s(5, 999999), s(5, 2), s(5, 10)}, []int{1, 2, 0}, 2},
		{"stable for equal timestamps", []time.Time{s(2, 0), s(1, 0), s(2, 0), s(1, 0)}, []int{1, 3, 0, 2}, 2},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			in := reorderTestFile(t, tt.ts...)
			ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
			if err != nil {
				t.Fatal(err)
			}
			var out bytes.Buffer
			n, err := ro.WriteTo(&out)
			if err != nil {
				t.Fatal(err)
			}
			if n != int64(len(in)) || out.Len() != len(in) {
				t.Errorf("wrote %d (%d) bytes, want %d", n, out.Len(), len(in))
			}
			if !bytes.Equal(out.Bytes()[:24], in[:24]) {
				t.Error("file header changed")
			}
			if got := packetIDs(t, out.Bytes()); !slices.Equal(got, tt.want) {
				t.Errorf("order %v, want %v", got, tt.want)
			}
			if ro.Packets() != len(tt.ts) || ro.OutOfOrder() != tt.outOfOrder || ro.Truncated() {
				t.Errorf("packets %d outOfOrder %d truncated %v", ro.Packets(), ro.OutOfOrder(), ro.Truncated())
			}
		})
	}
}

func TestReorderTruncated(t *testing.T) {
	in := reorderTestFile(t, time.Unix(2, 0), time.Unix(1, 0), time.Unix(0, 0))
	in = in[:len(in)-1]
	ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
	if err != nil {
		t.Fatal(err)
	}
	var out bytes.Buffer
	if _, err := ro.WriteTo(&out); err != nil {
		t.Fatal(err)
	}
	if !ro.Truncated() || !slices.Equal(packetIDs(t, out.Bytes()), []int{1, 0}) {
		t.Errorf("truncated %v, order %v", ro.Truncated(), packetIDs(t, out.Bytes()))
	}
}

func TestReorderErrors(t *testing.T) {
	for name, tc := range map[string]struct {
		in   []byte
		want string
	}{
		"no packets": {reorderTestFile(t), "no packets"},
		"garbage":    {bytes.Repeat([]byte("y"), 30), "unknown magic"},
	} {
		_, err := NewReorder(bytes.NewReader(tc.in), int64(len(tc.in)))
		if err == nil || !strings.Contains(err.Error(), tc.want) {
			t.Errorf("%s: error %v, want %q", name, err, tc.want)
		}
	}
}

// TestReorderManyPackets covers multiple parallel batches, including packets
// that are bigger than a whole batch, against a simple reference sort.
func TestReorderManyPackets(t *testing.T) {
	defer func(n int) { batchBytes = n }(batchBytes)
	batchBytes = 64 << 10

	r := rand.New(rand.NewPCG(1, 2))
	var buf bytes.Buffer
	w, err := NewWriter(&buf, LinkTypeEthernet)
	if err != nil {
		t.Fatal(err)
	}
	type pkt struct {
		ts   time.Time
		data []byte
	}
	var pkts []pkt
	for i := range 50000 {
		size := 8 + r.IntN(200)
		if i%10000 == 5 {
			size = batchBytes + 1000 // gets a batch by itself
		}
		data := make([]byte, size)
		binary.LittleEndian.PutUint64(data, uint64(i))
		p := pkt{time.Unix(int64(r.IntN(1000)), int64(r.IntN(1000))*1000), data}
		pkts = append(pkts, p)
		if err := w.WritePacket(p.ts, p.data); err != nil {
			t.Fatal(err)
		}
	}
	slices.SortStableFunc(pkts, func(a, b pkt) int { return a.ts.Compare(b.ts) })

	in := buf.Bytes()
	ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
	if err != nil {
		t.Fatal(err)
	}
	var out bytes.Buffer
	if _, err := ro.WriteTo(&out); err != nil {
		t.Fatal(err)
	}
	var want bytes.Buffer
	ww, _ := NewWriter(&want, LinkTypeEthernet)
	for _, p := range pkts {
		ww.WritePacket(p.ts, p.data)
	}
	if !bytes.Equal(out.Bytes(), want.Bytes()) {
		t.Error("reordered output differs from reference sort")
	}
}

type failingWriter struct{ n int }

func (f *failingWriter) Write(p []byte) (int, error) {
	if f.n--; f.n < 0 {
		return 0, errors.New("write failed")
	}
	return len(p), nil
}

func TestReorderWriteError(t *testing.T) {
	in := reorderTestFile(t, time.Unix(2, 0), time.Unix(1, 0))
	ro, err := NewReorder(bytes.NewReader(in), int64(len(in)))
	if err != nil {
		t.Fatal(err)
	}
	for n := range 2 {
		if _, err := ro.WriteTo(&failingWriter{n: n}); err == nil {
			t.Errorf("failing after %d writes: no error", n)
		}
	}
}