package rtph264

import (
	"bytes"
	"io"
	"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}, 512),
			),
		},
		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, 0x60, 0x44, 0xee, 0x88, 0x77, 0x79, 0xab,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x05,
				},
				[]byte{0x02, 0x03, 0x04, 0x05, 0x06, 0x07},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 181),
				[]byte{0x00, 0x01, 0x02, 0x03},
			),
			mergeBytes(
				[]byte{
					0x80, 0xe0, 0x44, 0xef, 0x88, 0x77, 0x79, 0xab,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x45,
				},
				[]byte{0x04, 0x05, 0x06, 0x07},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 147),
			),
		},
	},
	{
		"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 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 TestDecodePartOfFragmentedBeforeSingle(t *testing.T) {
	d := NewDecoder()

	_, _, err := d.Decode(mergeBytes(
		[]byte{
			0x80, 0xe0, 0x44, 0xef, 0x88, 0x77, 0x79, 0xab,
			0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x45,
		},
		[]byte{0x04, 0x05, 0x06, 0x07},
		bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 147),
	))
	require.Equal(t, ErrNonStartingPacketAndNoPrevious, err)

	_, _, err = d.Decode(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),
	))
	require.NoError(t, err)
}

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
	}{
		{
			"invalid rtp",
			[][]byte{
				{
					0xaa,
				},
			},
			"RTP header size insufficient: 1 < 4",
		},
		{
			"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 (invalid size)",
		},
		{
			"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 (invalid size)",
		},
		{
			"FU-A without payload",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeFUA),
			}},
			"invalid FU-A packet (invalid size)",
		},
		{
			"FU-A without start bit",
			[][]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),
				),
				{
					0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
					0x9d, 0xbb, 0x78, 0x12, byte(naluTypeFUA), 0x00,
				},
			},
			"invalid FU-A packet (non-starting)",
		},
		{
			"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 FU-A packet (invalid size)",
		},
		{
			"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,
					},
				),
			},
			"expected FU-A packet, got another type",
		},
		{
			"FU-A with two starting packets",
			[][]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, 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},
				),
			},
			"invalid FU-A packet (decoded two starting packets in a row)",
		},
		{
			"MTAP",
			[][]byte{{
				0x80, 0xe0, 0x44, 0xed, 0x88, 0x77, 0x6a, 0x15,
				0x9d, 0xbb, 0x78, 0x12, byte(naluTypeMTAP16),
			}},
			"packet type not supported (MTAP-16)",
		},
	} {
		t.Run(ca.name, func(t *testing.T) {
			d := NewDecoder()
			var err error
			for _, pkt := range ca.pkts {
				_, _, err = d.Decode(pkt)
			}
			require.Equal(t, ca.err, err.Error())
		})
	}
}

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 TestEncodeRandomInitialState(t *testing.T) {
	NewEncoder(96, nil, nil, nil)
}

type dummyReader struct {
	byts [][]byte
	i    int
}

func (f *dummyReader) Read(p []byte) (int, error) {
	if f.i >= len(f.byts) {
		return 0, io.EOF
	}
	n := copy(p, f.byts[f.i])
	f.i++
	return n, nil
}

func TestReadSPSPPS(t *testing.T) {
	for _, ca := range []struct {
		name string
		byts [][]byte
		sps  []byte
		pps  []byte
	}{
		{
			"sps then pps",
			[][]byte{
				{128, 96, 61, 205, 54, 67, 90, 125, 40, 249, 97, 176, 7, 1, 2},
				{128, 96, 61, 206, 54, 67, 90, 125, 40, 249, 97, 176, 8, 3, 4},
			},
			[]byte{0x07, 0x01, 0x02},
			[]byte{0x08, 0x03, 0x04},
		},
		{
			"pps then sps",
			[][]byte{
				{128, 96, 61, 206, 54, 67, 90, 125, 40, 249, 97, 176, 8, 3, 4},
				{128, 96, 61, 205, 54, 67, 90, 125, 40, 249, 97, 176, 7, 1, 2},
			},
			[]byte{0x07, 0x01, 0x02},
			[]byte{0x08, 0x03, 0x04},
		},
	} {
		t.Run(ca.name, func(t *testing.T) {
			sps, pps, err := NewDecoder().ReadSPSPPS(&dummyReader{byts: ca.byts})
			require.NoError(t, err)
			require.Equal(t, ca.sps, sps)
			require.Equal(t, ca.pps, pps)
		})
	}
}

func TestReadSPSPPSErrors(t *testing.T) {
	for _, ca := range []struct {
		name string
		byts [][]byte
		err  string
	}{
		{
			"empty",
			[][]byte{},
			"EOF",
		},
		{
			"more packets needed, then 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},
				),
			},
			"EOF",
		},
	} {
		t.Run(ca.name, func(t *testing.T) {
			_, _, err := NewDecoder().ReadSPSPPS(&dummyReader{byts: ca.byts})
			require.Equal(t, ca.err, err.Error())
		})
	}
}