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
}

type readerFunc func(p []byte) (int, error)

func (f readerFunc) Read(p []byte) (int, error) {
	return f(p)
}

var cases = []struct {
	name string
	dec  *NALUAndTimestamp
	enc  [][]byte
}{
	{
		"single",
		&NALUAndTimestamp{
			Timestamp: 25 * time.Millisecond,
			NALU: mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 8),
			),
		},
		[][]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",
		&NALUAndTimestamp{
			Timestamp: -20 * time.Millisecond,
			NALU: mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 8),
			),
		},
		[][]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",
		&NALUAndTimestamp{
			Timestamp: 55 * time.Millisecond,
			NALU: mergeBytes(
				[]byte{0x05},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 256),
			),
		},
		[][]byte{
			mergeBytes(
				[]byte{
					0x80, 0x60, 0x44, 0xed, 0x88, 0x77, 0x79, 0xab,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x85, 0x00, 0x01,
					0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 181),
				[]byte{0x00, 0x01},
			),
			mergeBytes(
				[]byte{
					0x80, 0xe0, 0x44, 0xee, 0x88, 0x77, 0x79, 0xab,
					0x9d, 0xbb, 0x78, 0x12, 0x1c, 0x45, 0x02, 0x03,
					0x04, 0x05, 0x06, 0x07,
				},
				bytes.Repeat([]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 73),
			),
		},
	},
}

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.dec)
			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) {
			i := 0
			r := readerFunc(func(p []byte) (int, error) {
				if i == len(ca.enc) {
					return 0, io.EOF
				}

				i++
				return copy(p, ca.enc[i-1]), nil
			})

			d := NewDecoder()

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

			dec, err := d.Read(r)
			require.NoError(t, err)
			require.Equal(t, ca.dec, dec)

			_, err = d.Read(r)
			require.Equal(t, io.EOF, err)
		})
	}
}

func TestDecodeAggregated(t *testing.T) {
	sent := false
	r := readerFunc(func(p []byte) (int, error) {
		if sent {
			return 0, io.EOF
		}

		sent = true
		pkt := []byte{
			0x80, 0xe0, 0x0e, 0x6a, 0x48, 0xf1, 0x7d, 0xb9,
			0x23, 0xe6, 0x5d, 0x50, 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,
		}
		return copy(p, pkt), nil
	})

	d := NewDecoder()

	nt, err := d.Read(r)
	require.NoError(t, err)
	require.Equal(t, &NALUAndTimestamp{
		NALU: []byte{0x09, 0xF0},
	}, nt)

	nt, err = d.Read(r)
	require.NoError(t, err)
	require.Equal(t, &NALUAndTimestamp{
		NALU: []byte{
			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,
		},
	}, nt)

	_, err = d.Read(r)
	require.Equal(t, io.EOF, err)
}