package media

import (
	"context"
	"net"
	"testing"
	"time"

	"github.com/pion/rtp"
)

func TestPacketGuardUsesPionAndAppliesCellPolicy(t *testing.T) {
	packet := &rtp.Packet{Header: rtp.Header{Version: 2, PayloadType: 8, SequenceNumber: 42, Timestamp: 160, SSRC: 99, Marker: true}, Payload: []byte{1, 2, 3, 4}}
	raw, err := packet.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	parsed, err := (PacketGuard{MaxPacketBytes: 1500, MaxPayloadBytes: 100, AllowedSSRC: 99, RequireSSRC: true, AllowedPayloadTypes: map[uint8]struct{}{8: {}}}).Parse(raw)
	if err != nil {
		t.Fatal(err)
	}
	if parsed.PayloadType != 8 || parsed.SequenceNumber != 42 || parsed.Timestamp != 160 || parsed.SSRC != 99 || !parsed.Marker || string(parsed.Payload) != string(packet.Payload) {
		t.Fatalf("unexpected parsed packet: %+v", parsed)
	}
}

func TestPacketGuardPreservesPCMAPayloadByteForByte(t *testing.T) {
	payload := make([]byte, 160)
	for index := range payload {
		payload[index] = byte((index*37 + 11) % 256)
	}
	packet := &rtp.Packet{Header: rtp.Header{Version: 2, PayloadType: 8, SequenceNumber: 700, Timestamp: 112000, SSRC: 0xA1B2C3D4}, Payload: payload}
	raw, err := packet.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	parsed, err := (PacketGuard{MaxPacketBytes: 1500, MaxPayloadBytes: 160, AllowedSSRC: 0xA1B2C3D4, RequireSSRC: true, AllowedPayloadTypes: map[uint8]struct{}{8: {}}}).Parse(raw)
	if err != nil {
		t.Fatal(err)
	}
	if parsed.PayloadType != 8 || parsed.SequenceNumber != packet.SequenceNumber || parsed.Timestamp != packet.Timestamp || parsed.SSRC != packet.SSRC {
		t.Fatalf("unexpected PCMA RTP header: %+v", parsed)
	}
	if string(parsed.Payload) != string(payload) {
		t.Fatalf("PCMA payload changed: got %x want %x", parsed.Payload, payload)
	}
}

func TestRTPStreamSendsAfterExplicitPeerBeforeInbound(t *testing.T) {
	receiver, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	if err != nil {
		t.Fatal(err)
	}
	defer receiver.Close()
	stream, err := ListenRTP("127.0.0.1:0", 118)
	if err != nil {
		t.Fatal(err)
	}
	defer stream.Close()
	if err := stream.SetPeer(receiver.LocalAddr().String()); err != nil {
		t.Fatal(err)
	}
	pcm := make([]byte, pcm16FrameSamples*2)
	if err := stream.SendPCM16(context.Background(), pcm, 16000); err != nil {
		t.Fatal(err)
	}
	if err := receiver.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, maxRTPPacketBytes)
	n, _, err := receiver.ReadFromUDP(buf)
	if err != nil {
		t.Fatal(err)
	}
	var packet rtp.Packet
	if err := packet.Unmarshal(buf[:n]); err != nil {
		t.Fatal(err)
	}
	if packet.PayloadType != 118 || len(packet.Payload) != len(pcm) {
		t.Fatalf("unexpected packet: payload_type=%d payload_bytes=%d", packet.PayloadType, len(packet.Payload))
	}
}

func TestALAWProfileNormalizesWireAudioToCanonicalPCM16(t *testing.T) {
	receiver, err := ListenRTPWithFormat("127.0.0.1:0", 8, FormatALAW, 8000)
	if err != nil {
		t.Fatal(err)
	}
	defer receiver.Close()
	sender, err := ListenRTPWithFormat("127.0.0.1:0", 8, FormatALAW, 8000)
	if err != nil {
		t.Fatal(err)
	}
	defer sender.Close()
	if err := sender.SetPeer(receiver.LocalAddr().String()); err != nil {
		t.Fatal(err)
	}
	pcm := make([]byte, pcm16FrameSamples*2)
	if err := sender.SendPCM16(context.Background(), pcm, 16000); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	decoded, err := receiver.ReadPayload(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if len(decoded) != len(pcm) {
		t.Fatalf("canonical PCM16 bytes=%d want=%d", len(decoded), len(pcm))
	}
}

func TestPacketGuardRejectsMalformedOrUnauthorizedMedia(t *testing.T) {
	guard := PacketGuard{MaxPacketBytes: 12, MaxPayloadBytes: 2, AllowedSSRC: 99, RequireSSRC: true, AllowedPayloadTypes: map[uint8]struct{}{8: {}}}
	if _, err := guard.Parse([]byte{0x80}); err == nil {
		t.Fatal("expected malformed RTP rejection")
	}
	packet := &rtp.Packet{Header: rtp.Header{Version: 2, PayloadType: 0, SequenceNumber: 1, Timestamp: 1, SSRC: 99}, Payload: []byte{1}}
	raw, err := packet.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := guard.Parse(raw); err == nil {
		t.Fatal("expected payload type rejection")
	}
	packet.PayloadType = 8
	packet.SSRC = 100
	raw, err = packet.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := guard.Parse(raw); err == nil {
		t.Fatal("expected SSRC rejection")
	}
	packet.SSRC = 99
	packet.Payload = []byte{1, 2, 3}
	raw, err = packet.Marshal()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := guard.Parse(raw); err == nil {
		t.Fatal("expected payload limit rejection")
	}
}