package device

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"testing"
)

type streamTestWriter struct{ bytes.Buffer }

func (*streamTestWriter) Close() error { return nil }

func TestPumpH264ForwardsConfigAndFrames(t *testing.T) {
	// The configuration payload deliberately resembles an invalid packet header.
	// Skipping it would corrupt alignment and lose all following video frames.
	config := []byte{0, 0, 0, 1, 0x67, 0x64, 0, 0x28, 0, 0, 0, 1, 0x68, 0xeb}
	keyframe := []byte{0, 0, 0, 1, 0x65, 1, 2, 3}
	next := []byte{0, 0, 0, 1, 0x41, 4, 5}
	var wire bytes.Buffer
	wire.WriteString("h264")
	// Initial session and later rotation both use standalone 12-byte headers.
	session := func(width, height uint32) {
		_ = binary.Write(&wire, binary.BigEndian, uint32(1<<31))
		_ = binary.Write(&wire, binary.BigEndian, width)
		_ = binary.Write(&wire, binary.BigEndian, height)
	}
	session(576, 1280)
	for i, payload := range [][]byte{config, keyframe, next} {
		if i == 2 {
			session(1280, 576)
		}
		var flags uint64
		if i == 0 {
			flags = 1 << 62
		}
		if i == 1 {
			flags = 1 << 61
		}
		_ = binary.Write(&wire, binary.BigEndian, flags)
		_ = binary.Write(&wire, binary.BigEndian, uint32(len(payload)))
		wire.Write(payload)
	}
	input, peer := net.Pipe()
	defer input.Close()
	go func() { _, _ = peer.Write(wire.Bytes()); _ = peer.Close() }()
	if err := readVideoCodec(input); err != nil {
		t.Fatal(err)
	}
	output := &streamTestWriter{}
	err := (&previewStream{}).pumpH264(input, output, make(chan struct{}))
	if err != io.EOF {
		t.Fatalf("unexpected end: %v", err)
	}
	want := bytes.Join([][]byte{config, keyframe, next}, nil)
	if !bytes.Equal(output.Bytes(), want) {
		t.Fatalf("lost codec config or packet alignment: got %x want %x", output.Bytes(), want)
	}
}

func TestPumpH264RejectsInvalidLength(t *testing.T) {
	for _, size := range []uint32{0, 16*1024*1024 + 1} {
		input, peer := net.Pipe()
		go func() {
			defer peer.Close()
			hdr := make([]byte, 12)
			binary.BigEndian.PutUint32(hdr[8:], size)
			_, _ = peer.Write(hdr)
		}()
		err := (&previewStream{}).pumpH264(input, &streamTestWriter{}, make(chan struct{}))
		input.Close()
		if err == nil || err == io.EOF {
			t.Fatalf("length %d was not rejected: %v", size, err)
		}
	}
}