package fmp4

import (
	"bytes"
	"testing"

	gomp4 "github.com/abema/go-mp4"
	"github.com/stretchr/testify/require"
)

func testMP4(t *testing.T, byts []byte, boxes []gomp4.BoxPath) {
	i := 0
	_, err := gomp4.ReadBoxStructure(bytes.NewReader(byts), func(h *gomp4.ReadHandle) (interface{}, error) {
		require.Equal(t, boxes[i], h.Path)
		i++
		return h.Expand()
	})
	require.NoError(t, err)
}

func TestPartMarshal(t *testing.T) {
	testVideoSamples := []*PartSample{
		{
			Duration: 2 * 90000,
			Payload: []byte{
				0x00, 0x00, 0x00, 0x04,
				0x01, 0x02, 0x03, 0x04, // SPS
				0x00, 0x00, 0x00, 0x01,
				0x08, // PPS
				0x00, 0x00, 0x00, 0x01,
				0x05, // IDR
			},
		},
		{
			Duration: 2 * 90000,
			Payload: []byte{
				0x00, 0x00, 0x00, 0x01,
				0x01, // non-IDR
			},
			IsNonSyncSample: true,
		},
		{
			Duration: 1 * 90000,
			Payload: []byte{
				0x00, 0x00, 0x00, 0x01,
				0x01, // non-IDR
			},
			IsNonSyncSample: true,
		},
	}

	testAudioSamples := []*PartSample{
		{
			Duration: 500 * 48000 / 1000,
			Payload: []byte{
				0x01, 0x02, 0x03, 0x04,
			},
		},
		{
			Duration: 1 * 48000,
			Payload: []byte{
				0x01, 0x02, 0x03, 0x04,
			},
		},
	}

	t.Run("video + audio", func(t *testing.T) {
		part := Part{
			Tracks: []*PartTrack{
				{
					ID:      1,
					Samples: testVideoSamples,
					IsVideo: true,
				},
				{
					ID:       2,
					BaseTime: 3 * 48000,
					Samples:  testAudioSamples,
				},
			},
		}

		byts, err := part.Marshal()
		require.NoError(t, err)

		boxes := []gomp4.BoxPath{
			{gomp4.BoxTypeMoof()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeMfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfdt()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTrun()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfdt()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTrun()},
			{gomp4.BoxTypeMdat()},
		}
		testMP4(t, byts, boxes)
	})

	t.Run("video only", func(t *testing.T) {
		part := Part{
			Tracks: []*PartTrack{
				{
					ID:      1,
					Samples: testVideoSamples,
					IsVideo: true,
				},
			},
		}

		byts, err := part.Marshal()
		require.NoError(t, err)

		boxes := []gomp4.BoxPath{
			{gomp4.BoxTypeMoof()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeMfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfdt()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTrun()},
			{gomp4.BoxTypeMdat()},
		}
		testMP4(t, byts, boxes)
	})

	t.Run("audio only", func(t *testing.T) {
		part := Part{
			Tracks: []*PartTrack{
				{
					ID:      1,
					Samples: testAudioSamples,
				},
			},
		}

		byts, err := part.Marshal()
		require.NoError(t, err)

		boxes := []gomp4.BoxPath{
			{gomp4.BoxTypeMoof()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeMfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfhd()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTfdt()},
			{gomp4.BoxTypeMoof(), gomp4.BoxTypeTraf(), gomp4.BoxTypeTrun()},
			{gomp4.BoxTypeMdat()},
		}
		testMP4(t, byts, boxes)
	})
}

func TestPartUnmarshal(t *testing.T) {
	byts := []byte{
		0x00, 0x00, 0x00, 0xd8, 0x6d, 0x6f, 0x6f, 0x66,
		0x00, 0x00, 0x00, 0x10, 0x6d, 0x66, 0x68, 0x64,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x70, 0x74, 0x72, 0x61, 0x66,
		0x00, 0x00, 0x00, 0x10, 0x74, 0x66, 0x68, 0x64,
		0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
		0x00, 0x00, 0x00, 0x14, 0x74, 0x66, 0x64, 0x74,
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x44,
		0x74, 0x72, 0x75, 0x6e, 0x01, 0x00, 0x0f, 0x01,
		0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0xe0,
		0x00, 0x02, 0xbf, 0x20, 0x00, 0x00, 0x00, 0x12,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x02, 0xbf, 0x20, 0x00, 0x00, 0x00, 0x05,
		0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x01, 0x5f, 0x90, 0x00, 0x00, 0x00, 0x05,
		0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x50, 0x74, 0x72, 0x61, 0x66,
		0x00, 0x00, 0x00, 0x10, 0x74, 0x66, 0x68, 0x64,
		0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
		0x00, 0x00, 0x00, 0x14, 0x74, 0x66, 0x64, 0x74,
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x02, 0x32, 0x80, 0x00, 0x00, 0x00, 0x24,
		0x74, 0x72, 0x75, 0x6e, 0x01, 0x00, 0x03, 0x01,
		0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0xfc,
		0x00, 0x00, 0x5d, 0xc0, 0x00, 0x00, 0x00, 0x04,
		0x00, 0x00, 0xbb, 0x80, 0x00, 0x00, 0x00, 0x04,
		0x00, 0x00, 0x00, 0x2c, 0x6d, 0x64, 0x61, 0x74,
		0x00, 0x00, 0x00, 0x04, 0x01, 0x02, 0x03, 0x04,
		0x00, 0x00, 0x00, 0x01, 0x08, 0x00, 0x00, 0x00,
		0x01, 0x05, 0x00, 0x00, 0x00, 0x01, 0x01, 0x00,
		0x00, 0x00, 0x01, 0x01, 0x01, 0x02, 0x03, 0x04,
		0x01, 0x02, 0x03, 0x04,
	}

	var parts Parts
	err := parts.Unmarshal(byts)
	require.NoError(t, err)

	require.Equal(t, Parts{{
		Tracks: []*PartTrack{
			{
				ID: 1,
				Samples: []*PartSample{
					{
						Duration: 2 * 90000,
						Payload: []byte{
							0x00, 0x00, 0x00, 0x04,
							0x01, 0x02, 0x03, 0x04, // SPS
							0x00, 0x00, 0x00, 0x01,
							0x08, // PPS
							0x00, 0x00, 0x00, 0x01,
							0x05, // IDR
						},
					},
					{
						Duration: 2 * 90000,
						Payload: []byte{
							0x00, 0x00, 0x00, 0x01,
							0x01, // non-IDR
						},
						IsNonSyncSample: true,
					},
					{
						Duration: 1 * 90000,
						Payload: []byte{
							0x00, 0x00, 0x00, 0x01,
							0x01, // non-IDR
						},
						IsNonSyncSample: true,
					},
				},
			},
			{
				ID:       2,
				BaseTime: 3 * 48000,
				Samples: []*PartSample{
					{
						Duration: 500 * 48000 / 1000,
						Payload: []byte{
							0x01, 0x02, 0x03, 0x04,
						},
					},
					{
						Duration: 1 * 48000,
						Payload: []byte{
							0x01, 0x02, 0x03, 0x04,
						},
					},
				},
			},
		},
	}}, parts)
}