package h264

import (
	"testing"

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

var casesSPS = []struct {
	name   string
	byts   []byte
	sps    SPS
	width  int
	height int
	fps    float64
}{
	{
		"352x288",
		[]byte{
			0x67, 0x64, 0x00, 0x0c, 0xac, 0x3b, 0x50, 0xb0,
			0x4b, 0x42, 0x00, 0x00, 0x03, 0x00, 0x02, 0x00,
			0x00, 0x03, 0x00, 0x3d, 0x08,
		},
		SPS{
			ProfileIdc:                     100,
			LevelIdc:                       12,
			ChromaFormatIdc:                1,
			Log2MaxFrameNumMinus4:          6,
			PicOrderCntType:                2,
			MaxNumRefFrames:                1,
			GapsInFrameNumValueAllowedFlag: true,
			PicWidthInMbsMinus1:            21,
			PicHeightInMapUnitsMinus1:      17,
			FrameMbsOnlyFlag:               true,
			Direct8x8InferenceFlag:         true,
			VUI: &SPS_VUI{
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1,
					TimeScale:          30,
					FixedFrameRateFlag: true,
				},
			},
		},
		352,
		288,
		15,
	},
	{
		"1280x720",
		[]byte{
			0x67, 0x64, 0x00, 0x1f, 0xac, 0xd9, 0x40, 0x50,
			0x05, 0xbb, 0x01, 0x6c, 0x80, 0x00, 0x00, 0x03,
			0x00, 0x80, 0x00, 0x00, 0x1e, 0x07, 0x8c, 0x18,
			0xcb,
		},
		SPS{
			ProfileIdc:                  100,
			LevelIdc:                    31,
			ChromaFormatIdc:             1,
			Log2MaxPicOrderCntLsbMinus4: 2,
			MaxNumRefFrames:             4,
			PicWidthInMbsMinus1:         79,
			PicHeightInMapUnitsMinus1:   44,
			FrameMbsOnlyFlag:            true,
			Direct8x8InferenceFlag:      true,
			VUI: &SPS_VUI{
				AspectRatioInfoPresentFlag: true,
				AspectRatioIdc:             1,
				VideoSignalTypePresentFlag: true,
				VideoFormat:                5,
				VideoFullRangeFlag:         true,
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick: 1,
					TimeScale:      60,
				},
				BitstreamRestriction: &SPS_BitstreamRestriction{
					MotionVectorsOverPicBoundariesFlag: true,
					Log2MaxMvLengthHorizontal:          11,
					Log2MaxMvLengthVertical:            11,
					MaxNumReorderFrames:                2,
					MaxDecFrameBuffering:               4,
				},
			},
		},
		1280,
		720,
		30,
	},
	{
		"1920x1080 baseline",
		[]byte{
			0x67, 0x42, 0xc0, 0x28, 0xd9, 0x00, 0x78, 0x02,
			0x27, 0xe5, 0x84, 0x00, 0x00, 0x03, 0x00, 0x04,
			0x00, 0x00, 0x03, 0x00, 0xf0, 0x3c, 0x60, 0xc9, 0x20,
		},
		SPS{
			ProfileIdc:                66,
			ChromaFormatIdc:           1,
			ConstraintSet0Flag:        true,
			ConstraintSet1Flag:        true,
			LevelIdc:                  40,
			PicOrderCntType:           2,
			MaxNumRefFrames:           3,
			PicWidthInMbsMinus1:       119,
			PicHeightInMapUnitsMinus1: 67,
			FrameMbsOnlyFlag:          true,
			Direct8x8InferenceFlag:    true,
			FrameCropping: &SPS_FrameCropping{
				BottomOffset: 4,
			},
			VUI: &SPS_VUI{
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick: 1,
					TimeScale:      60,
				},
				BitstreamRestriction: &SPS_BitstreamRestriction{
					MotionVectorsOverPicBoundariesFlag: true,
					Log2MaxMvLengthHorizontal:          11,
					Log2MaxMvLengthVertical:            11,
					MaxDecFrameBuffering:               3,
				},
			},
		},
		1920,
		1080,
		30,
	},
	{
		"1920x1080 nvidia",
		[]byte{
			0x67, 0x64, 0x00, 0x28, 0xac, 0xd9, 0x40, 0x78,
			0x02, 0x27, 0xe5, 0x84, 0x00, 0x00, 0x03, 0x00,
			0x04, 0x00, 0x00, 0x03, 0x00, 0xf0, 0x3c, 0x60,
			0xc6, 0x58,
		},
		SPS{
			ProfileIdc:                  100,
			LevelIdc:                    40,
			ChromaFormatIdc:             1,
			Log2MaxPicOrderCntLsbMinus4: 2,
			MaxNumRefFrames:             4,
			PicWidthInMbsMinus1:         119,
			PicHeightInMapUnitsMinus1:   67,
			FrameMbsOnlyFlag:            true,
			Direct8x8InferenceFlag:      true,
			FrameCropping: &SPS_FrameCropping{
				BottomOffset: 4,
			},
			VUI: &SPS_VUI{
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick: 1,
					TimeScale:      60,
				},
				BitstreamRestriction: &SPS_BitstreamRestriction{
					MotionVectorsOverPicBoundariesFlag: true,
					Log2MaxMvLengthHorizontal:          11,
					Log2MaxMvLengthVertical:            11,
					MaxNumReorderFrames:                2,
					MaxDecFrameBuffering:               4,
				},
			},
		},
		1920,
		1080,
		30,
	},
	{
		"1920x1080",
		[]byte{
			0x67, 0x64, 0x00, 0x29, 0xac, 0x13, 0x31, 0x40,
			0x78, 0x04, 0x47, 0xde, 0x03, 0xea, 0x02, 0x02,
			0x03, 0xe0, 0x00, 0x00, 0x03, 0x00, 0x20, 0x00,
			0x00, 0x06, 0x52, // 0x80,
		},
		SPS{
			ProfileIdc:                  100,
			LevelIdc:                    41,
			ChromaFormatIdc:             1,
			Log2MaxFrameNumMinus4:       8,
			Log2MaxPicOrderCntLsbMinus4: 5,
			MaxNumRefFrames:             4,
			PicWidthInMbsMinus1:         119,
			PicHeightInMapUnitsMinus1:   33,
			Direct8x8InferenceFlag:      true,
			FrameCropping: &SPS_FrameCropping{
				BottomOffset: 2,
			},
			VUI: &SPS_VUI{
				AspectRatioInfoPresentFlag:   true,
				AspectRatioIdc:               1,
				OverscanInfoPresentFlag:      true,
				OverscanAppropriateFlag:      true,
				VideoSignalTypePresentFlag:   true,
				VideoFormat:                  5,
				ColourDescriptionPresentFlag: true,
				ColourPrimaries:              1,
				TransferCharacteristics:      1,
				MatrixCoefficients:           1,
				ChromaLocInfoPresentFlag:     true,
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1,
					TimeScale:          50,
					FixedFrameRateFlag: true,
				},
				PicStructPresentFlag: true,
			},
		},
		1920,
		1080,
		25,
	},
	{
		"hikvision",
		[]byte{103, 100, 0, 32, 172, 23, 42, 1, 64, 30, 104, 64, 0, 1, 194, 0, 0, 87, 228, 33},
		SPS{
			ProfileIdc:                  100,
			LevelIdc:                    32,
			ChromaFormatIdc:             1,
			Log2MaxPicOrderCntLsbMinus4: 4,
			MaxNumRefFrames:             1,
			PicWidthInMbsMinus1:         79,
			PicHeightInMapUnitsMinus1:   59,
			FrameMbsOnlyFlag:            true,
			Direct8x8InferenceFlag:      true,
			Log2MaxFrameNumMinus4:       10,
			VUI: &SPS_VUI{
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1800,
					TimeScale:          90000,
					FixedFrameRateFlag: true,
				},
			},
		},
		1280,
		960,
		25,
	},
	{
		"scaling matrix",
		[]byte{
			103, 100, 0, 50, 173, 132, 1, 12, 32, 8, 97, 0, 67, 8, 2,
			24, 64, 16, 194, 0, 132, 59, 80, 20, 0, 90, 211,
			112, 16, 16, 20, 0, 0, 3, 0, 4, 0, 0, 3, 0, 162, 16,
		},
		SPS{
			ProfileIdc:      100,
			LevelIdc:        50,
			ChromaFormatIdc: 1,
			ScalingList4x4: [][]int32{
				{
					16, 16, 16, 16, 16, 16, 16, 16,
					16, 16, 16, 16, 16, 16, 16, 16,
				},
				{
					16, 16, 16, 16, 16, 16, 16, 16,
					16, 16, 16, 16, 16, 16, 16, 16,
				},
				{
					16, 16, 16, 16, 16, 16, 16, 16,
					16, 16, 16, 16, 16, 16, 16, 16,
				},
				{
					16, 16, 16, 16, 16, 16, 16, 16,
					16, 16, 16, 16, 16, 16, 16, 16,
				},
				{
					16, 16, 16, 16, 16, 16, 16, 16,
					16, 16, 16, 16, 16, 16, 16, 16,
				},
				{
					16, 16, 16, 16, 16, 16, 16, 16,
					16, 16, 16, 16, 16, 16, 16, 16,
				},
			},
			UseDefaultScalingMatrix4x4Flag: []bool{
				false, false, false, false, false, false,
			},
			Log2MaxFrameNumMinus4:          6,
			PicOrderCntType:                2,
			MaxNumRefFrames:                1,
			GapsInFrameNumValueAllowedFlag: true,
			PicWidthInMbsMinus1:            159,
			PicHeightInMapUnitsMinus1:      89,
			FrameMbsOnlyFlag:               true,
			Direct8x8InferenceFlag:         true,
			VUI: &SPS_VUI{
				VideoSignalTypePresentFlag:   true,
				VideoFormat:                  5,
				VideoFullRangeFlag:           true,
				ColourDescriptionPresentFlag: true,
				ColourPrimaries:              1,
				TransferCharacteristics:      1,
				MatrixCoefficients:           1,
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1,
					TimeScale:          40,
					FixedFrameRateFlag: true,
				},
			},
		},
		2560,
		1440,
		20,
	},
	{
		"1920x1080 nvenc hrd",
		[]byte{
			103, 100, 0, 42, 172, 44, 172, 7,
			128, 34, 126, 92, 5, 168, 8, 8,
			10, 0, 0, 7, 208, 0, 3, 169,
			129, 192, 0, 0, 76, 75, 0, 0,
			38, 37, 173, 222, 92, 20,
		},
		SPS{
			ProfileIdc:                  100,
			LevelIdc:                    42,
			ChromaFormatIdc:             1,
			Log2MaxFrameNumMinus4:       4,
			Log2MaxPicOrderCntLsbMinus4: 4,
			MaxNumRefFrames:             2,
			PicWidthInMbsMinus1:         119,
			PicHeightInMapUnitsMinus1:   67,
			FrameMbsOnlyFlag:            true,
			Direct8x8InferenceFlag:      true,
			FrameCropping: &SPS_FrameCropping{
				BottomOffset: 4,
			},
			VUI: &SPS_VUI{
				AspectRatioInfoPresentFlag:   true,
				AspectRatioIdc:               1,
				VideoSignalTypePresentFlag:   true,
				VideoFormat:                  5,
				ColourDescriptionPresentFlag: true,
				ColourPrimaries:              1,
				TransferCharacteristics:      1,
				MatrixCoefficients:           1,
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1000,
					TimeScale:          120000,
					FixedFrameRateFlag: true,
				},
				NalHRD: &SPS_HRD{
					BitRateValueMinus1:                 []uint32{39061},
					CpbSizeValueMinus1:                 []uint32{156249},
					CbrFlag:                            []bool{true},
					InitialCpbRemovalDelayLengthMinus1: 23,
					CpbRemovalDelayLengthMinus1:        15,
					DpbOutputDelayLengthMinus1:         5,
					TimeOffsetLength:                   24,
				},
				PicStructPresentFlag: true,
			},
		},
		1920,
		1080,
		60,
	},
	{
		"1920x1080 hikvision nal hrd + vcl hrd",
		[]byte{
			103, 77, 0, 41, 154, 100, 3, 192,
			17, 63, 46, 2, 220, 4, 4, 5,
			0, 0, 3, 3, 232, 0, 0, 195,
			80, 232, 96, 0, 186, 180, 0, 2,
			234, 196, 187, 203, 141, 12, 0, 23,
			86, 128, 0, 93, 88, 151, 121, 112,
			160,
		},
		SPS{
			ProfileIdc:                  77,
			ChromaFormatIdc:             1,
			LevelIdc:                    41,
			Log2MaxFrameNumMinus4:       5,
			Log2MaxPicOrderCntLsbMinus4: 5,
			MaxNumRefFrames:             1,
			PicWidthInMbsMinus1:         119,
			PicHeightInMapUnitsMinus1:   67,
			FrameMbsOnlyFlag:            true,
			Direct8x8InferenceFlag:      true,
			FrameCropping: &SPS_FrameCropping{
				BottomOffset: 4,
			},
			VUI: &SPS_VUI{
				AspectRatioInfoPresentFlag:   true,
				AspectRatioIdc:               1,
				VideoSignalTypePresentFlag:   true,
				VideoFormat:                  5,
				VideoFullRangeFlag:           true,
				ColourDescriptionPresentFlag: true,
				ColourPrimaries:              1,
				TransferCharacteristics:      1,
				MatrixCoefficients:           1,
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1000,
					TimeScale:          50000,
					FixedFrameRateFlag: true,
				},
				NalHRD: &SPS_HRD{
					BitRateScale:                       4,
					CpbSizeScale:                       3,
					BitRateValueMinus1:                 []uint32{11948},
					CpbSizeValueMinus1:                 []uint32{95585},
					CbrFlag:                            []bool{false},
					InitialCpbRemovalDelayLengthMinus1: 23,
					CpbRemovalDelayLengthMinus1:        15,
					DpbOutputDelayLengthMinus1:         5,
					TimeOffsetLength:                   24,
				},
				VclHRD: &SPS_HRD{
					BitRateScale:                       4,
					CpbSizeScale:                       3,
					BitRateValueMinus1:                 []uint32{11948},
					CpbSizeValueMinus1:                 []uint32{95585},
					CbrFlag:                            []bool{false},
					InitialCpbRemovalDelayLengthMinus1: 23,
					CpbRemovalDelayLengthMinus1:        15,
					DpbOutputDelayLengthMinus1:         5,
					TimeOffsetLength:                   24,
				},
				PicStructPresentFlag: true,
			},
		},
		1920,
		1080,
		25,
	},
	{
		"1920x1080 mbs_only_flag = 0",
		[]byte{
			0x67, 0x4d, 0x40, 0x28, 0xab, 0x60, 0x3c, 0x02,
			0x23, 0xef, 0x01, 0x10, 0x00, 0x00, 0x03, 0x00,
			0x10, 0x00, 0x00, 0x03, 0x03, 0x2e, 0x94, 0x00,
			0x35, 0x64, 0x06, 0xb2, 0x85, 0x08, 0x0e, 0xe2,
			0xc5, 0x22, 0xc0,
		},
		SPS{
			ProfileIdc:                  77,
			ChromaFormatIdc:             1,
			ConstraintSet1Flag:          true,
			LevelIdc:                    40,
			Log2MaxFrameNumMinus4:       1,
			Log2MaxPicOrderCntLsbMinus4: 2,
			MaxNumRefFrames:             2,
			PicWidthInMbsMinus1:         119,
			PicHeightInMapUnitsMinus1:   33,
			Direct8x8InferenceFlag:      true,
			FrameCropping: &SPS_FrameCropping{
				BottomOffset: 2,
			},
			VUI: &SPS_VUI{
				AspectRatioInfoPresentFlag: true,
				AspectRatioIdc:             1,
				TimingInfo: &SPS_TimingInfo{
					NumUnitsInTick:     1,
					TimeScale:          50,
					FixedFrameRateFlag: true,
				},
				NalHRD: &SPS_HRD{
					BitRateScale:                       4,
					CpbSizeScale:                       10,
					BitRateValueMinus1:                 []uint32{3416},
					CpbSizeValueMinus1:                 []uint32{213},
					CbrFlag:                            []bool{false},
					InitialCpbRemovalDelayLengthMinus1: 20,
					CpbRemovalDelayLengthMinus1:        5,
					DpbOutputDelayLengthMinus1:         1,
				},
				PicStructPresentFlag: true,
				BitstreamRestriction: &SPS_BitstreamRestriction{
					MotionVectorsOverPicBoundariesFlag: true,
					MaxBytesPerPicDenom:                2,
					Log2MaxMvLengthHorizontal:          10,
					Log2MaxMvLengthVertical:            9,
					MaxNumReorderFrames:                1,
					MaxDecFrameBuffering:               4,
				},
			},
		},
		1920,
		1080,
		25,
	},
}

func TestSPSUnmarshal(t *testing.T) {
	for _, ca := range casesSPS {
		t.Run(ca.name, func(t *testing.T) {
			var sps SPS
			err := sps.Unmarshal(ca.byts)
			require.NoError(t, err)
			require.Equal(t, ca.sps, sps)
			require.Equal(t, ca.width, sps.Width())
			require.Equal(t, ca.height, sps.Height())
			require.Equal(t, ca.fps, sps.FPS())
		})
	}
}

func BenchmarkSPSUnmarshal(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var sps SPS
		sps.Unmarshal([]byte{ //nolint:errcheck
			103, 77, 0, 41, 154, 100, 3, 192,
			17, 63, 46, 2, 220, 4, 4, 5,
			0, 0, 3, 3, 232, 0, 0, 195,
			80, 232, 96, 0, 186, 180, 0, 2,
			234, 196, 187, 203, 141, 12, 0, 23,
			86, 128, 0, 93, 88, 151, 121, 112,
			160,
		})
	}
}

func FuzzSPSUnmarshal(f *testing.F) {
	for _, ca := range casesSPS {
		f.Add(ca.byts)
	}

	f.Fuzz(func(_ *testing.T, b []byte) {
		var sps SPS
		err := sps.Unmarshal(b)
		if err != nil {
			return
		}

		sps.Width()
		sps.Height()
		sps.FPS()
	})
}
