package rtph264

import (
	"bytes"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func mergeBytes(vals ...[]byte) []byte {
	size := 0
	for _, v := range vals {
		size += len(v)
	}
	res := make([]byte, size)

	pos := 0
	for _, v := range vals {
		n := copy(res[pos:], v)
		pos += n
	}

	return res
}

var cases = []struct {
	name  string
	nalus [][]byte
	pts   time.Duration
	enc   [][]byte
}{
	{
		"single",
		[][]byte{
			mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 8),
			),
		},
		25 * time.Millisecond,
		[][]byte{
			mergeBytes(
				[]byte{
					0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6f, 0x1f,
					0x9d, 0xbb, 0x78, 0x12, 0x05,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 8),
			),
		},
	},
	{
		"negative timestamp",
		[][]byte{
			mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 8),
			),
		},
		-20 * time.Millisecond,
		[][]byte{
			mergeBytes(
				[]byte{
					0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x5f, 0x4d,
					0x9d, 0xbb, 0x78, 0x12, 0x05,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 8),
			),
		},
	},
	{
		"fragmented",
		[][]byte{
			mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 256),
			),
		},
		55 * time.Millisecond,
		[][]byte{
			mergeBytes(
				[]byte{
					0x80, 0x60, 0x44, 0xed, 0x88, 0x77, 0x79, 0xab,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x85,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 182),
				[]byte{0x00, 0x01},
			),
			mergeBytes(
				[]byte{
					0x80, 0xe0, 0x44, 0xee, 0x88, 0x77, 0x79, 0xab,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x45,
				},
				[]byte{0x02, 0x03, 0x04, 0x05, 0x06, 0x07},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 73),
			),
		},
	},
	{
		"aggregated",
		[][]byte{
			{0x09, 0xF0},
			{
				0x41, 0x9a, 0x24, 0x6c, 0x41, 0x4f, 0xfe, 0xd6,
				0x8c, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x6d, 0x40,
			},
		},
		0,
		[][]byte{
			{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x66, 0x55,
				0x9d, 0xbb, 0x78, 0x12, 0x18, 0x00, 0x02, 0x09,
				0xf0, 0x00, 0x44, 0x41, 0x9a, 0x24, 0x6c, 0x41,
				0x4f, 0xfe, 0xd6, 0x8c, 0xb0, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x6d, 0x40,
			},
		},
	},
	{
		"aggregated followed by single",
		[][]byte{
			{0x09, 0xF0},
			{
				0x41, 0x9a, 0x24, 0x6c, 0x41, 0x4f, 0xfe, 0xd6,
				0x8c, 0xb0, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x6d, 0x40,
			},
			mergeBytes(
				[]byte{0x08},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 175),
			),
		},
		0,
		[][]byte{
			{
				0x80, 0x60, 0x44, 0xed, 0x88, 0x77, 0x66, 0x55,
				0x9d, 0xbb, 0x78, 0x12, 0x18, 0x00, 0x02, 0x09,
				0xf0, 0x00, 0x44, 0x41, 0x9a, 0x24, 0x6c, 0x41,
				0x4f, 0xfe, 0xd6, 0x8c, 0xb0, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00,
				0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03, 0x00,
				0x00, 0x03, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
				0x00, 0x00, 0x03, 0x00, 0x00, 0x6d, 0x40,
			},
			mergeBytes(
				[]byte{
					0x80, 0xe0, 0x44, 0xee, 0x88, 0x77, 0x66, 0x55,
					0x9d, 0xbb, 0x78, 0x12, 0x08,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 175),
			),
		},
	},
	{
		"fragmented followed by aggregated",
		[][]byte{
			mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 256),
			),
			{0x09, 0xF0},
			{0x09, 0xF0},
		},
		0,
		[][]byte{
			mergeBytes(
				[]byte{
					0x80, 0x60, 0x44, 0xed, 0x88, 0x77, 0x66, 0x55,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x85,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 182),
				[]byte{0x00, 0x01},
			),
			mergeBytes(
				[]byte{
					0x80, 0x60, 0x44, 0xee, 0x88, 0x77, 0x66, 0x55,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x45,
				},
				[]byte{0x02, 0x03, 0x04, 0x05, 0x06, 0x07},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 73),
			),
			{
				0x80, 0xe0, 0x44, 0xef, 0x88, 0x77, 0x66, 0x55,
				0x9d, 0xbb, 0x78, 0x12, 0x18, 0x00, 0x02, 0x09,
				0xf0, 0x00, 0x02, 0x09, 0xf0,
			},
		},
	},
}

func TestEncode(t *testing.T) {
	for _, ca := range cases {
		t.Run(ca.name, func(t *testing.T) {
			sequenceNumber := uint16(0x44ed)
			ssrc := uint32(0x9dbb7812)
			initialTs := uint32(0x88776655)
			e := NewEncoder(96, &sequenceNumber, &ssrc, &initialTs)
			enc, err := e.Encode(ca.nalus, ca.pts)
			require.NoError(t, err)
			require.Equal(t, ca.enc, enc)
		})
	}
}

func TestDecode(t *testing.T) {
	for _, ca := range cases {
		t.Run(ca.name, func(t *testing.T) {
			d := NewDecoder()

			// send an initial packet downstream
			// in order to compute the timestamp,
			// which is relative to the initial packet
			_, _, err := d.Decode([]byte{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x66, 0x55,
				0x9d, 0xbb, 0x78, 0x12, 0x06, 0x00,
			})
			require.NoError(t, err)

			var nalus [][]byte

			for _, pkt := range ca.enc {
				addNALUs, pts, err := d.Decode(pkt)
				if err == ErrMorePacketsNeeded {
					continue
				}

				require.NoError(t, err)
				require.Equal(t, ca.pts, pts)
				nalus = append(nalus, addNALUs...)
			}

			require.Equal(t, ca.nalus, nalus)
		})
	}
}

func TestDecodeSTAPAWithPadding(t *testing.T) {
	d := NewDecoder()
	nalus, _, err := d.Decode([]byte{
		0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x66, 0x55,
		0x9d, 0xbb, 0x78, 0x12, 0x18, 0x00, 0x02, 0xaa,
		0xbb, 0x00, 0x02, 0xcc, 0xdd, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	})
	require.NoError(t, err)
	require.Equal(t, [][]byte{
		{0xaa, 0xbb},
		{0xcc, 0xdd},
	}, nalus)
}

func TestDecodeErrors(t *testing.T) {
	for _, ca := range []struct {
		name string
		pkts [][]byte
		err  string
	}{
		{
			"missing payload",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12,
			}},
			"payload is too short",
		},
		{
			"STAP-A without NALUs",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeSTAPA),
			}},
			"STAP-A packet doesn't contain any NALU",
		},
		{
			"STAP-A without size",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeSTAPA), 0x01,
			}},
			"Invalid STAP-A packet",
		},
		{
			"STAP-A with invalid size",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeSTAPA), 0x00, 0x15,
			}},
			"Invalid STAP-A packet",
		},
		{
			"FU-A without payload",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeFUA),
			}},
			"Invalid FU-A packet",
		},
		{
			"FU-A without start bit",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeFUA), 0x00,
			}},
			"first NALU does not contain the start bit",
		},
		{
			"FU-A with 2nd packet empty",
			[][]byte{
				mergeBytes(
					[]byte{
						0x80, 0x60, 0x44, 0xed, 0x88, 0x77, 0x79, 0xab,
						0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x85,
					},
					bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 182),
					[]byte{0x00, 0x01},
				),
				mergeBytes(
					[]byte{
						0x80, 0xe0, 0x44, 0xee, 0x88, 0x77, 0x79, 0xab,
						0x9d, 0xbb, 0x78, 0x12,
					},
				),
			},
			"Invalid non-starting FU-A packet",
		},
		{
			"FU-A with 2nd packet invalid",
			[][]byte{
				mergeBytes(
					[]byte{
						0x80, 0x60, 0x44, 0xed, 0x88, 0x77, 0x79, 0xab,
						0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x85,
					},
					bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 182),
					[]byte{0x00, 0x01},
				),
				mergeBytes(
					[]byte{
						0x80, 0xe0, 0x44, 0xee, 0x88, 0x77, 0x79, 0xab,
						0x9d, 0xbb, 0x78, 0x12, 0x01, 0x00,
					},
				),
			},
			"Packet is not FU-A",
		},
		{
			"MTAP",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeMTAP16),
			}},
			"NALU type not supported (MTAP16)",
		},
	} {
		t.Run(ca.name, func(t *testing.T) {
			d := NewDecoder()
			var err error
			for _, pkt := range ca.pkts {
				_, _, err = d.Decode(pkt)
			}
			require.NotEqual(t, ErrMorePacketsNeeded, err)
			require.Equal(t, ca.err, err.Error())
		})
	}
}