package mts

import (
	"bytes"
	"math/rand"
	"testing"
	"time"

	"bitbucket.org/ausocean/av/container/mts/pes"
	"bitbucket.org/ausocean/av/container/mts/psi"
)

func TestGetPTSRange(t *testing.T) {
	const (
		numOfFrames = 20
		maxFrame    = 1000
		minFrame    = 100
		rate        = 25 // fps
		interval    = float64(1) / rate
		ptsFreq     = 90000 // Hz
	)

	// Generate randomly sized data for each frame.
	rand.Seed(time.Now().UnixNano())
	frames := make([][]byte, numOfFrames)
	for i := range frames {
		size := rand.Intn(maxFrame-minFrame) + minFrame
		frames[i] = make([]byte, size)
	}

	var clip bytes.Buffer

	// Write PAT and PMT.
	pat := Packet{
		PUSI:    true,
		PID:     PatPid,
		CC:      0,
		AFC:     HasPayload,
		Payload: psi.AddPadding(patTable),
	}
	_, err := clip.Write(pat.Bytes(nil))
	if err != nil {
		t.Fatalf("did not expect error writing PAT: %v", err)
	}

	// Create mts packet from pmt table.
	pmt := Packet{
		PUSI:    true,
		PID:     PmtPid,
		CC:      0,
		AFC:     HasPayload,
		Payload: psi.AddPadding(pmtTable),
	}
	_, err = clip.Write(pmt.Bytes(nil))
	if err != nil {
		t.Fatalf("did not expected error writing PMT: %v", err)
	}

	var curTime float64
	for _, frame := range frames {
		nextPTS := curTime * ptsFreq
		// Prepare PES data.
		pesPkt := pes.Packet{
			StreamID:     videoStreamID,
			PDI:          hasPTS,
			PTS:          uint64(nextPTS),
			Data:         frame,
			HeaderLength: 5,
		}

		buf := pesPkt.Bytes(nil)

		pusi := true
		for len(buf) != 0 {
			pkt := Packet{
				PUSI: pusi,
				PID:  videoPid,
				RAI:  pusi,
				CC:   0,
				AFC:  hasAdaptationField | hasPayload,
				PCRF: pusi,
			}
			n := pkt.FillPayload(buf)
			buf = buf[n:]

			pusi = false
			_, err := clip.Write(pkt.Bytes(nil))
			if err != nil {
				t.Fatalf("did not expect error writing video packet: %v", err)
			}
		}
		curTime += interval
	}

	got, err := GetPTSRange(clip.Bytes())
	if err != nil {
		t.Fatalf("did not expect error getting PTS range: %v", err)
	}

	want := [2]uint64{0, uint64((numOfFrames - 1) * interval * ptsFreq)}
	if got != want {
		t.Errorf("did not get expected result.\n Got: %v\n Want: %v\n", got, want)
	}
}