// Copyright 2023 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package buffer

import (
	"encoding/binary"
	"fmt"
	"math"
	"sync"
	"testing"
	"time"

	"github.com/pion/rtcp"
	"github.com/pion/rtp"
	"github.com/pion/sdp/v3"
	"github.com/pion/webrtc/v4"
	"github.com/stretchr/testify/require"

	"github.com/livekit/mediatransportutil/pkg/nack"

	act "github.com/livekit/livekit-server/pkg/sfu/rtpextension/abscapturetime"
	"github.com/livekit/livekit-server/pkg/sfu/utils"
)

var h265Codec = webrtc.RTPCodecParameters{
	RTPCodecCapability: webrtc.RTPCodecCapability{
		MimeType:  "video/h265",
		ClockRate: 90000,
		RTCPFeedback: []webrtc.RTCPFeedback{{
			Type: "nack",
		}},
	},
	PayloadType: 116,
}

var vp8Codec = webrtc.RTPCodecParameters{
	RTPCodecCapability: webrtc.RTPCodecCapability{
		MimeType:  "video/vp8",
		ClockRate: 90000,
		RTCPFeedback: []webrtc.RTCPFeedback{{
			Type: "nack",
		}},
	},
	PayloadType: 96,
}

var opusCodec = webrtc.RTPCodecParameters{
	RTPCodecCapability: webrtc.RTPCodecCapability{
		MimeType:  "audio/opus",
		ClockRate: 48000,
	},
	PayloadType: 111,
}

func TestNack(t *testing.T) {
	t.Run("nack normal", func(t *testing.T) {
		buff := NewBuffer(123, 1, 1)
		buff.codecType = webrtc.RTPCodecTypeVideo
		require.NotNil(t, buff)
		var wg sync.WaitGroup
		// 5 tries
		wg.Add(5)
		buff.OnRtcpFeedback(func(fb []rtcp.Packet) {
			for _, pkt := range fb {
				switch p := pkt.(type) {
				case *rtcp.TransportLayerNack:
					if p.Nacks[0].PacketList()[0] == 1 && p.MediaSSRC == 123 {
						wg.Done()
					}
				}
			}
		})
		buff.Bind(webrtc.RTPParameters{
			HeaderExtensions: nil,
			Codecs:           []webrtc.RTPCodecParameters{vp8Codec},
		}, vp8Codec.RTPCodecCapability, 0)
		rtt := uint32(20)
		buff.nacker.SetRTT(rtt)
		for i := range 15 {
			if i == 1 {
				continue
			}
			if i < 14 {
				time.Sleep(time.Duration(float64(rtt)*math.Pow(nack.NackQueueParamsDefault.BackoffFactor, float64(i))+10) * time.Millisecond)
			} else {
				time.Sleep(500 * time.Millisecond) // even a long wait should not exceed max retries
			}
			pkt := rtp.Packet{
				Header: rtp.Header{
					Version:        2,
					PayloadType:    96,
					SequenceNumber: uint16(i),
					Timestamp:      uint32(i),
					SSRC:           123,
				},
				Payload: []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1},
			}
			b, err := pkt.Marshal()
			require.NoError(t, err)
			_, err = buff.Write(b)
			require.NoError(t, err)
		}
		wg.Wait()

	})

	t.Run("nack with seq wrap", func(t *testing.T) {
		buff := NewBuffer(123, 1, 1)
		buff.codecType = webrtc.RTPCodecTypeVideo
		require.NotNil(t, buff)
		var wg sync.WaitGroup
		expects := map[uint16]int{
			65534: 0,
			65535: 0,
			0:     0,
			1:     0,
		}
		wg.Add(5 * len(expects)) // retry 5 times
		buff.OnRtcpFeedback(func(fb []rtcp.Packet) {
			for _, pkt := range fb {
				switch p := pkt.(type) {
				case *rtcp.TransportLayerNack:
					if p.MediaSSRC == 123 {
						for _, v := range p.Nacks {
							v.Range(func(seq uint16) bool {
								if _, ok := expects[seq]; ok {
									wg.Done()
								} else {
									require.Fail(t, "unexpected nack seq ", seq)
								}
								return true
							})
						}
					}
				}
			}
		})
		buff.Bind(webrtc.RTPParameters{
			HeaderExtensions: nil,
			Codecs:           []webrtc.RTPCodecParameters{vp8Codec},
		}, vp8Codec.RTPCodecCapability, 0)
		rtt := uint32(30)
		buff.nacker.SetRTT(rtt)
		for i := range 15 {
			if i > 0 && i < 5 {
				continue
			}
			if i < 14 {
				time.Sleep(time.Duration(float64(rtt)*math.Pow(nack.NackQueueParamsDefault.BackoffFactor, float64(i))+10) * time.Millisecond)
			} else {
				time.Sleep(500 * time.Millisecond) // even a long wait should not exceed max retries
			}
			pkt := rtp.Packet{
				Header: rtp.Header{
					Version:        2,
					PayloadType:    96,
					SequenceNumber: uint16(i + 65533),
					Timestamp:      uint32(i),
					SSRC:           123,
				},
				Payload: []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1},
			}
			b, err := pkt.Marshal()
			require.NoError(t, err)
			_, err = buff.Write(b)
			require.NoError(t, err)
		}
		wg.Wait()

	})
}

func TestNackOnRTX(t *testing.T) {
	buff := NewBuffer(123, 1, 1)
	buff.codecType = webrtc.RTPCodecTypeVideo
	var mu sync.Mutex
	var nacked []uint16
	buff.OnRtcpFeedback(func(fb []rtcp.Packet) {
		mu.Lock()
		defer mu.Unlock()
		for _, pkt := range fb {
			if p, ok := pkt.(*rtcp.TransportLayerNack); ok && p.MediaSSRC == 123 {
				for _, pair := range p.Nacks {
					nacked = append(nacked, pair.PacketList()...)
				}
			}
		}
	})
	require.NoError(t, buff.Bind(webrtc.RTPParameters{
		Codecs: []webrtc.RTPCodecParameters{vp8Codec},
	}, vp8Codec.RTPCodecCapability, 0))
	buff.rtxPayloadType = 97
	rtt := uint32(20)
	buff.nacker.SetRTT(rtt)

	payload := []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1}
	// sequence numbers 1 and 2 are lost
	for _, sn := range []uint16{0, 3} {
		pkt := rtp.Packet{
			Header:  rtp.Header{Version: 2, PayloadType: 96, SequenceNumber: sn, Timestamp: uint32(sn), SSRC: 123},
			Payload: payload,
		}
		b, err := pkt.Marshal()
		require.NoError(t, err)
		_, err = buff.Write(b)
		require.NoError(t, err)
	}

	time.Sleep(time.Duration(2*rtt) * time.Millisecond)
	mu.Lock()
	nacked = nil
	mu.Unlock()

	// RTX repairs sequence number 2, the NACK for 1 is due and must be sent
	rtxPkt := rtp.Packet{
		Header:  rtp.Header{Version: 2, PayloadType: 97, SequenceNumber: 500, Timestamp: 2, SSRC: 456},
		Payload: append([]byte{0x00, 0x02}, payload...),
	}
	buff.writeRTX(&rtxPkt, time.Now().UnixNano())

	mu.Lock()
	defer mu.Unlock()
	require.Equal(t, []uint16{1}, nacked)
}

func TestNewBuffer(t *testing.T) {
	tests := []struct {
		name string
	}{
		{
			name: "Must not be nil and add packets in sequence",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var TestPackets = []*rtp.Packet{
				{
					Header: rtp.Header{
						Version:        2,
						PayloadType:    96,
						SequenceNumber: 65533,
						SSRC:           123,
					},
				},
				{
					Header: rtp.Header{
						Version:        2,
						PayloadType:    96,
						SequenceNumber: 65534,
						SSRC:           123,
					},
					Payload: []byte{1},
				},
				{
					Header: rtp.Header{
						Version:        2,
						PayloadType:    96,
						SequenceNumber: 2,
						SSRC:           123,
					},
				},
				{
					Header: rtp.Header{
						Version:        2,
						PayloadType:    96,
						SequenceNumber: 65535,
						SSRC:           123,
					},
				},
			}
			buff := NewBuffer(123, 1, 1)
			buff.codecType = webrtc.RTPCodecTypeVideo
			require.NotNil(t, buff)
			buff.OnRtcpFeedback(func(_ []rtcp.Packet) {})
			buff.Bind(webrtc.RTPParameters{
				HeaderExtensions: nil,
				Codecs:           []webrtc.RTPCodecParameters{vp8Codec},
			}, vp8Codec.RTPCodecCapability, 0)

			for _, p := range TestPackets {
				buf, _ := p.Marshal()
				_, _ = buff.Write(buf)
			}
			require.Equal(t, uint16(2), buff.rtpStats.HighestSequenceNumber())
			require.Equal(t, uint64(65536+2), buff.rtpStats.ExtendedHighestSequenceNumber())
		})
	}
}

func TestFractionLostReport(t *testing.T) {
	buff := NewBuffer(123, 1, 1)
	require.NotNil(t, buff)

	var wg sync.WaitGroup

	// with loss proxying
	wg.Add(1)
	buff.SetAudioLossProxying(true)
	buff.SetLastFractionLostReport(55)
	buff.OnRtcpFeedback(func(fb []rtcp.Packet) {
		for _, pkt := range fb {
			switch p := pkt.(type) {
			case *rtcp.ReceiverReport:
				for _, v := range p.Reports {
					require.EqualValues(t, 55, v.FractionLost)
				}
				wg.Done()
			}
		}
	})
	buff.Bind(webrtc.RTPParameters{
		HeaderExtensions: nil,
		Codecs:           []webrtc.RTPCodecParameters{opusCodec},
	}, opusCodec.RTPCodecCapability, 0)
	for i := range 15 {
		pkt := rtp.Packet{
			Header: rtp.Header{
				Version:        2,
				PayloadType:    111,
				SequenceNumber: uint16(i),
				Timestamp:      uint32(i),
				SSRC:           123,
			},
			Payload: []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1},
		}
		b, err := pkt.Marshal()
		require.NoError(t, err)
		if i == 1 {
			time.Sleep(1 * time.Second)
		}
		_, err = buff.Write(b)
		require.NoError(t, err)
	}
	wg.Wait()

	wg.Add(1)
	buff.SetAudioLossProxying(false)
	buff.OnRtcpFeedback(func(fb []rtcp.Packet) {
		for _, pkt := range fb {
			switch p := pkt.(type) {
			case *rtcp.ReceiverReport:
				for _, v := range p.Reports {
					require.EqualValues(t, 0, v.FractionLost)
				}
				wg.Done()
			}
		}
	})
	buff.Bind(webrtc.RTPParameters{
		HeaderExtensions: nil,
		Codecs:           []webrtc.RTPCodecParameters{opusCodec},
	}, opusCodec.RTPCodecCapability, 0)
	for i := range 15 {
		pkt := rtp.Packet{
			Header: rtp.Header{
				Version:        2,
				PayloadType:    111,
				SequenceNumber: uint16(i),
				Timestamp:      uint32(i),
				SSRC:           123,
			},
			Payload: []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1},
		}
		b, err := pkt.Marshal()
		require.NoError(t, err)
		if i == 1 {
			time.Sleep(1 * time.Second)
		}
		_, err = buff.Write(b)
		require.NoError(t, err)
	}
	wg.Wait()
}

func TestCodecChange(t *testing.T) {
	// codec change before bind
	buff := NewBuffer(123, 1, 1)
	require.NotNil(t, buff)
	changedCodec := make(chan webrtc.RTPCodecParameters, 1)
	buff.OnCodecChange(func(rp webrtc.RTPCodecParameters) {
		select {
		case changedCodec <- rp:
		default:
			t.Fatalf("codec change not consumed")
		}
	})
	buff.OnStreamRestart(func(reason string) {
		require.Equal(t, "codec-change", reason)

		// read once to clear pending restart
		var buf [1500]byte
		extPkt, err := buff.ReadExtended(buf[:])
		require.NoError(t, err)
		require.Nil(t, extPkt)
	})

	h265Pkt := rtp.Packet{
		Header: rtp.Header{
			Version:        2,
			PayloadType:    116,
			SequenceNumber: 1,
			Timestamp:      1,
			SSRC:           123,
		},
		Payload: []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1},
	}
	buf, err := h265Pkt.Marshal()
	require.NoError(t, err)
	_, err = buff.Write(buf)
	require.NoError(t, err)

	select {
	case <-changedCodec:
		t.Fatalf("unexpected codec change")
	case <-time.After(100 * time.Millisecond):
	}

	// Bind sets up VP8 as expected codec,
	// packet written to the buffer above before Bind is H.265,
	// that should trigger a codec change and a stream restart
	// when the queued packets from Write before Bind get flushed
	buff.Bind(
		webrtc.RTPParameters{
			HeaderExtensions: nil,
			Codecs:           []webrtc.RTPCodecParameters{vp8Codec, h265Codec},
		},
		vp8Codec.RTPCodecCapability,
		0,
	)

	select {
	case c := <-changedCodec:
		require.Equal(t, h265Codec, c)
	case <-time.After(1 * time.Second):
		t.Fatalf("expected codec change")
	}

	// second codec change - writing VP8 packet after Bind should trigger another codec change
	vp8Pkt := rtp.Packet{
		Header: rtp.Header{
			Version:        2,
			PayloadType:    96,
			SequenceNumber: 3,
			Timestamp:      3,
			SSRC:           123,
		},
		Payload: []byte{0xff, 0xff, 0xff, 0xfd, 0xb4, 0x9f, 0x94, 0x1},
	}
	buf, err = vp8Pkt.Marshal()
	require.NoError(t, err)
	_, err = buff.Write(buf)
	require.NoError(t, err)

	select {
	case c := <-changedCodec:
		require.Equal(t, vp8Codec, c)
	case <-time.After(1 * time.Second):
		t.Fatalf("expected codec change")
	}
	fmt.Printf("done second codec change\n") // REMOVE

	// out of order pkts can't cause codec change
	// rewrite the VP8 packet to start the sequence after a stream restart
	_, err = buff.Write(buf)
	require.NoError(t, err)

	h265Pkt.SequenceNumber = 2
	h265Pkt.Timestamp = 2
	buf, err = h265Pkt.Marshal()
	require.NoError(t, err)
	_, err = buff.Write(buf)
	require.NoError(t, err)
	select {
	case <-changedCodec:
		t.Fatalf("unexpected codec change")
	case <-time.After(100 * time.Millisecond):
	}

	// unknown codec should not cause change even if it is in order
	h265Pkt.SequenceNumber = 4
	h265Pkt.Timestamp = 4
	h265Pkt.PayloadType = 117
	buf, err = h265Pkt.Marshal()
	require.NoError(t, err)
	_, err = buff.Write(buf)
	require.NoError(t, err)
	select {
	case <-changedCodec:
		t.Fatalf("unexpected codec change")
	case <-time.After(100 * time.Millisecond):
	}

	// an in-order packet should change codec again
	h265Pkt.SequenceNumber = 5
	h265Pkt.Timestamp = 5
	h265Pkt.PayloadType = 116
	buf, err = h265Pkt.Marshal()
	require.NoError(t, err)
	_, err = buff.Write(buf)
	require.NoError(t, err)
	select {
	case c := <-changedCodec:
		require.Equal(t, h265Codec, c)
	case <-time.After(1 * time.Second):
		t.Fatalf("expected codec change")
	}
}

// Write unmarshals into reused scratch and callers reuse the input buffer, so each ExtPacket
// must keep its own header and extensions, and Write plus ReadExtended must not allocate per packet
func TestIngestNoAlloc(t *testing.T) {
	testCases := []struct {
		name      string
		codec     webrtc.RTPCodecParameters
		payload   []byte
		maxAllocs float64
	}{
		{"opus", opusCodec, make([]byte, 100), 0},
		// VP8 boxes the parsed descriptor into ExtPacket.Payload
		{"vp8", vp8Codec, []byte{0x90, 0xe0, 0x80, 0x01, 0x05, 0x00, 0x01}, 1},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			buff := NewBuffer(123, 100, 100)
			require.NoError(t, buff.Bind(webrtc.RTPParameters{
				HeaderExtensions: []webrtc.RTPHeaderExtensionParameter{
					{URI: sdp.AudioLevelURI, ID: 1},
					{URI: act.AbsCaptureTimeURI, ID: 2},
				},
				Codecs: []webrtc.RTPCodecParameters{tc.codec},
			}, tc.codec.RTPCodecCapability, 0))
			defer buff.Close()

			marshal := func(sn uint16, level byte) []byte {
				p := rtp.Packet{
					Header: rtp.Header{
						Version:        2,
						PayloadType:    uint8(tc.codec.PayloadType),
						SequenceNumber: sn,
						Timestamp:      uint32(sn) * 960,
						SSRC:           123,
					},
					Payload: tc.payload,
				}
				require.NoError(t, p.SetExtension(1, []byte{level}))
				require.NoError(t, p.SetExtension(3, []byte{0, 0}))
				raw, err := p.Marshal()
				require.NoError(t, err)
				return raw
			}

			// one input buffer for all writes, like the transport read loop
			input := make([]byte, 1500)
			write := func(sn uint16, level byte) {
				n := copy(input, marshal(sn, level))
				_, err := buff.Write(input[:n])
				require.NoError(t, err)
			}

			readBuf := make([]byte, 1500)
			write(1, 10)
			write(2, 20)
			for _, want := range []struct {
				sn    uint16
				level byte
			}{{1, 10}, {2, 20}} {
				ep, err := buff.ReadExtended(readBuf)
				require.NoError(t, err)
				require.Equal(t, want.sn, ep.Packet.SequenceNumber)
				require.Equal(t, []byte{want.level}, ep.Packet.GetExtension(1))
				require.Equal(t, tc.payload, ep.Packet.Payload)
				require.Nil(t, ep.AbsCaptureTimeExt)
				ReleaseExtPacket(ep)
			}

			raw := marshal(3, 30)
			sn := uint16(3)
			ingest := func() {
				binary.BigEndian.PutUint16(raw[2:], sn)
				binary.BigEndian.PutUint32(raw[4:], uint32(sn)*960)
				sn++
				_, err := buff.Write(raw)
				require.NoError(t, err)
				ep, err := buff.ReadExtended(readBuf)
				require.NoError(t, err)
				ReleaseExtPacket(ep)
			}
			allocs := testing.AllocsPerRun(1000, ingest)
			if !utils.RaceEnabled {
				require.LessOrEqual(t, allocs, tc.maxAllocs, "allocations per Write plus ReadExtended")
			}
		})
	}
}

// a dropped padding only packet shifts later sequence numbers, ReadExtended must keep the shifted number
// in the header and in RawPacket
func TestIngestDroppedPaddingKeepsAdjustedSequenceNumber(t *testing.T) {
	buff := NewBuffer(123, 100, 100)
	require.NoError(t, buff.Bind(webrtc.RTPParameters{
		Codecs: []webrtc.RTPCodecParameters{opusCodec},
	}, opusCodec.RTPCodecCapability, 0))
	defer buff.Close()

	write := func(sn uint16, paddingOnly bool) {
		p := rtp.Packet{
			Header: rtp.Header{
				Version:        2,
				PayloadType:    uint8(opusCodec.PayloadType),
				SequenceNumber: sn,
				Timestamp:      uint32(sn) * 960,
				SSRC:           123,
			},
		}
		if paddingOnly {
			p.Header.Padding = true
			p.Header.PaddingSize = 20
		} else {
			p.Payload = make([]byte, 100)
		}
		raw, err := p.Marshal()
		require.NoError(t, err)
		_, err = buff.Write(raw)
		require.NoError(t, err)
	}

	write(40, false)
	write(41, true)
	write(42, false)

	readBuf := make([]byte, 1500)
	for _, wantSN := range []uint16{40, 41} {
		ep, err := buff.ReadExtended(readBuf)
		require.NoError(t, err)
		require.Equal(t, wantSN, ep.Packet.SequenceNumber)
		require.Equal(t, wantSN, binary.BigEndian.Uint16(ep.RawPacket[2:]))
		ReleaseExtPacket(ep)
	}
}

func BenchmarkMemcpu(b *testing.B) {
	buf := make([]byte, 1500*1500*10)
	buf2 := make([]byte, 1500*1500*20)

	for b.Loop() {
		copy(buf2, buf)
	}
}

func BenchmarkExtPacketFactory(b *testing.B) {

	for b.Loop() {
		extPkt := ExtPacketFactory.Get().(*ExtPacket)
		*extPkt = ExtPacket{}
		ExtPacketFactory.Put(extPkt)
	}
}