package h264

import (
	"bytes"
	"os"
	"path/filepath"
	"testing"
)

// streamNALs builds an Annex-B byte stream from a list of NAL unit types.
// Each NAL gets a minimal one-byte payload so the type is easy to assert.
func streamNALs(types ...byte) []byte {
	var out []byte
	for _, t := range types {
		out = append(out, 0x00, 0x00, 0x01) // 3-byte start code
		out = append(out, 0x60|t)           // nal_ref_idc=3, type in low bits
		out = append(out, 0x00)
	}
	return out
}

// firstNALType extracts the type of the first NAL in a byte slice.
func firstNALType(au []byte) byte {
	if len(au) < 4 {
		return 0
	}
	// Skip the start code: 3 or 4 bytes.
	pay := 3
	if au[0] == 0 && au[1] == 0 && au[2] == 0 {
		pay = 4
	}
	return au[pay] & 0x1F
}

func typesOf(aus [][]byte) [][]byte {
	var types [][]byte
	for _, au := range aus {
		var ts []byte
		rest := au
		for len(rest) > 0 {
			ts = append(ts, firstNALType(rest))
			// advance past this NAL (start code + payload of >=1 byte)
			skip := 4
			if rest[0] == 0 && rest[1] == 0 && rest[2] == 0 {
				skip = 5
			}
			// find next start code
			next := bytes.Index(rest[skip:], []byte{0x00, 0x00, 0x01})
			if next < 0 {
				break
			}
			rest = rest[skip+next:]
		}
		types = append(types, ts)
	}
	return types
}

func TestSplitAUsEmpty(t *testing.T) {
	if got := SplitAUs(nil); got != nil {
		t.Errorf("SplitAUs(nil) = %v, want nil", got)
	}
}

func TestSplitAUsNoStartCode(t *testing.T) {
	data := []byte{0x67, 0x42, 0x00}
	got := SplitAUs(data)
	if len(got) != 1 || !bytes.Equal(got[0], data) {
		t.Errorf("SplitAUs passthrough = %v, want [% x]", got, data)
	}
}

func TestSplitAUsKeyframeThenP(t *testing.T) {
	// SPS(7), PPS(8), IDR(5), P(1), P(1)
	stream := streamNALs(7, 8, 5, 1, 1)
	aus := SplitAUs(stream)
	if len(aus) != 3 {
		t.Fatalf("got %d AUs, want 3", len(aus))
	}
	types := typesOf(aus)
	want := [][]byte{{7, 8, 5}, {1}, {1}}
	for i := range want {
		if !bytes.Equal(types[i], want[i]) {
			t.Errorf("AU %d types = %v, want %v", i, types[i], want[i])
		}
	}
}

func TestSplitAUsWithAUD(t *testing.T) {
	// AUD(9), IDR(5), P(1)
	stream := streamNALs(9, 5, 1)
	aus := SplitAUs(stream)
	if len(aus) != 2 {
		t.Fatalf("got %d AUs, want 2", len(aus))
	}
	types := typesOf(aus)
	want := [][]byte{{9, 5}, {1}}
	for i := range want {
		if !bytes.Equal(types[i], want[i]) {
			t.Errorf("AU %d types = %v, want %v", i, types[i], want[i])
		}
	}
}

func TestSplitAUsAUDThenPrefix(t *testing.T) {
	// AUD(9), SPS(7), PPS(8), IDR(5), P(1)
	stream := streamNALs(9, 7, 8, 5, 1)
	aus := SplitAUs(stream)
	if len(aus) != 2 {
		t.Fatalf("got %d AUs, want 2", len(aus))
	}
	types := typesOf(aus)
	want := [][]byte{{9, 7, 8, 5}, {1}}
	for i := range want {
		if !bytes.Equal(types[i], want[i]) {
			t.Errorf("AU %d types = %v, want %v", i, types[i], want[i])
		}
	}
}

func TestSplitAUsFourByteCodes(t *testing.T) {
	var stream []byte
	for _, t := range []byte{7, 8, 5, 1} {
		stream = append(stream, 0x00, 0x00, 0x00, 0x01) // 4-byte
		stream = append(stream, 0x60|t, 0x00)
	}
	aus := SplitAUs(stream)
	if len(aus) != 2 {
		t.Fatalf("got %d AUs, want 2", len(aus))
	}
}

func TestSplitAUsMixedCodes(t *testing.T) {
	var stream []byte
	stream = append(stream, 0x00, 0x00, 0x01, 0x67, 0x00)       // 3-byte SPS
	stream = append(stream, 0x00, 0x00, 0x00, 0x01, 0x68, 0x00) // 4-byte PPS
	stream = append(stream, 0x00, 0x00, 0x00, 0x01, 0x65, 0x00) // 4-byte IDR
	stream = append(stream, 0x00, 0x00, 0x01, 0x41, 0x00)       // 3-byte P

	aus := SplitAUs(stream)
	if len(aus) != 2 {
		t.Fatalf("got %d AUs, want 2", len(aus))
	}
	types := typesOf(aus)
	want := [][]byte{{7, 8, 5}, {1}}
	for i := range want {
		if !bytes.Equal(types[i], want[i]) {
			t.Errorf("AU %d types = %v, want %v", i, types[i], want[i])
		}
	}
}

func TestSplitAUsReassemblesStream(t *testing.T) {
	stream := streamNALs(7, 8, 5, 1, 1, 7, 8, 5, 1)
	aus := SplitAUs(stream)
	var rebuilt []byte
	for _, au := range aus {
		rebuilt = append(rebuilt, au...)
	}
	if !bytes.Equal(rebuilt, stream) {
		t.Error("reassembled stream differs from input")
	}
}

// TestSplitAUsRealFile validates against a real libx264 Annex-B stream.
func TestSplitAUsRealFile(t *testing.T) {
	path := filepath.Join("..", "..", "testdata", "sample.h264")
	data, err := os.ReadFile(path)
	if err != nil {
		t.Skipf("testdata not available: %v", err)
	}

	aus := SplitAUs(data)
	if len(aus) < 2 {
		t.Fatalf("got %d AUs, want >= 2", len(aus))
	}

	// First AU must be a keyframe: SPS + PPS + optional prefix NALs (SEI)
	// followed by an IDR slice.
	first := typesOf(aus[:1])[0]
	if len(first) < 3 || first[0] != 7 || first[1] != 8 || first[len(first)-1] != 5 {
		t.Errorf("first AU types = %v, want prefix [7 8 ... 5]", first)
	}

	// All AUs must be non-empty and start with a valid start code.
	for i, au := range aus {
		if len(au) < 4 {
			t.Fatalf("AU %d too short: %d bytes", i, len(au))
		}
		if !(au[0] == 0 && au[1] == 0 && (au[2] == 0 || au[2] == 1)) {
			t.Errorf("AU %d does not start with a start code: % x", i, au[:4])
		}
	}

	// Reassembly must reproduce the input exactly.
	var rebuilt []byte
	for _, au := range aus {
		rebuilt = append(rebuilt, au...)
	}
	if !bytes.Equal(rebuilt, data) {
		t.Error("reassembled stream differs from input")
	}
}