package pcap

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

// testFile builds a pcap file with one packet per entry of sizes.
func testFile(t *testing.T, sizes ...int) []byte {
	t.Helper()
	var buf bytes.Buffer
	w, err := NewWriter(&buf, LinkTypeEthernet)
	if err != nil {
		t.Fatal(err)
	}
	for i, n := range sizes {
		if err := w.WritePacket(time.Unix(int64(i), 0), bytes.Repeat([]byte{byte(i)}, n)); err != nil {
			t.Fatal(err)
		}
	}
	return buf.Bytes()
}

// splitParts runs Split and returns the parts' contents.
func splitParts(t *testing.T, in []byte, limit SplitLimit) ([][]byte, SplitStats, error) {
	t.Helper()
	var parts []*bytes.Buffer
	st, err := Split(bytes.NewReader(in), limit, func(part int) (io.Writer, error) {
		if part != len(parts)+1 {
			t.Errorf("create(%d), want %d", part, len(parts)+1)
		}
		parts = append(parts, &bytes.Buffer{})
		return parts[len(parts)-1], nil
	})
	out := make([][]byte, len(parts))
	for i, p := range parts {
		out[i] = p.Bytes()
	}
	return out, st, err
}

// packetSizes returns the captured lengths of the packets in a pcap file.
func packetSizes(t *testing.T, file []byte) []int {
	t.Helper()
	r, err := NewReader(bytes.NewReader(file))
	if err != nil {
		t.Fatal(err)
	}
	var sizes []int
	for {
		rec, err := r.Next()
		if err == io.EOF {
			return sizes
		}
		if err != nil {
			t.Fatal(err)
		}
		sizes = append(sizes, len(rec)-16)
	}
}

func TestSplit(t *testing.T) {
	tests := []struct {
		name  string
		sizes []int
		limit SplitLimit
		want  [][]int // packet sizes per part
	}{
		{"by packets", []int{10, 20, 30, 40, 50}, SplitLimit{Packets: 2}, [][]int{{10, 20}, {30, 40}, {50}}},
		{"by packets exact", []int{10, 20, 30, 40}, SplitLimit{Packets: 2}, [][]int{{10, 20}, {30, 40}}},
		{"one part", []int{10, 20}, SplitLimit{Packets: 5}, [][]int{{10, 20}}},
		// 24 byte header + 2*(16+20) = 96 fits exactly; a third packet does not.
		{"by bytes", []int{20, 20, 20, 20, 20}, SplitLimit{Bytes: 96}, [][]int{{20, 20}, {20, 20}, {20}}},
		{"oversized packet gets its own part", []int{10, 500, 10}, SplitLimit{Bytes: 100}, [][]int{{10}, {500}, {10}}},
		{"both limits", []int{1, 1, 1, 100, 1}, SplitLimit{Packets: 2, Bytes: 100}, [][]int{{1, 1}, {1}, {100}, {1}}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			in := testFile(t, tt.sizes...)
			parts, st, err := splitParts(t, in, tt.limit)
			if err != nil {
				t.Fatal(err)
			}
			if st.Parts != len(tt.want) || st.Packets != len(tt.sizes) || st.Truncated {
				t.Errorf("stats %+v", st)
			}
			if len(parts) != len(tt.want) {
				t.Fatalf("got %d parts, want %d", len(parts), len(tt.want))
			}
			var joined []byte
			for i, p := range parts {
				if !bytes.Equal(p[:24], in[:24]) {
					t.Errorf("part %d: file header %x, want %x", i+1, p[:24], in[:24])
				}
				if got := packetSizes(t, p); !slices.Equal(got, tt.want[i]) {
					t.Errorf("part %d: packet sizes %v, want %v", i+1, got, tt.want[i])
				}
				if tt.limit.Bytes > 0 && len(tt.want[i]) > 1 && int64(len(p)) > tt.limit.Bytes {
					t.Errorf("part %d: %d bytes exceeds limit %d", i+1, len(p), tt.limit.Bytes)
				}
				joined = append(joined, p[24:]...)
			}
			if !bytes.Equal(joined, in[24:]) {
				t.Error("parts do not add up to the input records")
			}
		})
	}
}

func TestSplitTruncated(t *testing.T) {
	in := testFile(t, 10, 20, 30)
	parts, st, err := splitParts(t, in[:len(in)-5], SplitLimit{Packets: 10})
	if err != nil {
		t.Fatal(err)
	}
	if !st.Truncated || st.Packets != 2 || len(parts) != 1 {
		t.Errorf("stats %+v", st)
	}
}

func TestSplitErrors(t *testing.T) {
	pcapng := []byte{0x0a, 0x0d, 0x0d, 0x0a, 0x1c, 0, 0, 0, 0x4d, 0x3c, 0x2b, 0x1a, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
	corrupt := testFile(t, 10)
	binary.LittleEndian.PutUint32(corrupt[24+8:], MaxRecordLen+1)
	for name, tc := range map[string]struct {
		in   []byte
		want string
	}{
		"empty input": {nil, "shorter than"},
		"no packets":  {testFile(t), "no packets"},
		"pcapng":      {pcapng, "pcapng"},
		"garbage":     {bytes.Repeat([]byte("x"), 40), "unknown magic"},
		"corrupt":     {corrupt, "corrupt"},
	} {
		_, _, err := splitParts(t, tc.in, SplitLimit{Packets: 1})
		if err == nil || !strings.Contains(err.Error(), tc.want) {
			t.Errorf("%s: error %v, want %q", name, err, tc.want)
		}
	}
}

func TestSplitCreateError(t *testing.T) {
	boom := errors.New("boom")
	_, err := Split(bytes.NewReader(testFile(t, 1, 1)), SplitLimit{Packets: 1}, func(part int) (io.Writer, error) {
		if part == 2 {
			return nil, boom
		}
		return io.Discard, nil
	})
	if !errors.Is(err, boom) {
		t.Errorf("err %v, want %v", err, boom)
	}
}

func TestReaderByteOrderAndNanoseconds(t *testing.T) {
	// Big endian, nanosecond resolution, link type 101, one 2 byte packet.
	in := []byte{
		0xa1, 0xb2, 0x3c, 0x4d, 0, 2, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 0, 0, 0, 101,
		0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 2, 0xaa, 0xbb,
	}
	r, err := NewReader(bytes.NewReader(in))
	if err != nil {
		t.Fatal(err)
	}
	if r.LinkType() != LinkTypeRaw {
		t.Errorf("link type %d", r.LinkType())
	}
	rec, err := r.Next()
	if err != nil || !bytes.Equal(rec, in[24:]) {
		t.Errorf("record %x, %v", rec, err)
	}
	if _, err := r.Next(); err != io.EOF {
		t.Errorf("err %v, want EOF", err)
	}
}