seek-tune/signal/webrtc.go

// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
// SPDX-License-Identifier: MIT

//go:build !js
// +build !js

// save-to-disk is a simple application that shows how to record your webcam/microphone using Pion WebRTC and save VP8/Opus to disk.
package signal

import (
	"fmt"
	"io"
	"time"

	"github.com/pion/interceptor"
	"github.com/pion/interceptor/pkg/intervalpli"
	"github.com/pion/webrtc/v4"
	"go.mongodb.org/mongo-driver/bson/primitive"

	"song-recognition/shazam"

	"github.com/pion/webrtc/v4/pkg/media"
)

func SaveToDisk(i media.Writer, track *webrtc.TrackRemote) {
	defer func() {
		if err := i.Close(); err != nil {
			panic(err)
		}
	}()

	for {
		rtpPacket, _, err := track.ReadRTP()
		if err != nil {
			fmt.Println(err)
			return
		}
		if err := i.WriteRTP(rtpPacket); err != nil {
			fmt.Println(err)
			return
		}
	}
}

func SaveToBytes(track *webrtc.TrackRemote) ([]byte, error) {
	var audioData []byte

	for {
		rtpPacket, _, err := track.ReadRTP()
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}

		// Extract audio payload from RTP packet
		payload := rtpPacket.Payload

		// Append audio payload to audioData
		audioData = append(audioData, payload...)
		// fmt.Println("ByteArray: ", audioData)
	}

	return audioData, nil
}

func MatchSampleAudio(track *webrtc.TrackRemote) ([]primitive.M, error) {
	// Use time.After to stop after 15 seconds
	stop := time.After(50 * time.Second)

	// Use a ticker to process sampleAudio every 2 seconds
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	var sampleAudio []byte
	var matches []primitive.M

	for {
		select {
		case <-ticker.C:
			// Process sampleAudio every 2 seconds
			if len(sampleAudio) > 0 {
				matchess, err := shazam.Match(sampleAudio)
				matches = matchess
				if err != nil {
					fmt.Println(err)
					return nil, nil
				}

				// Reset sampleAudio for fresh input
				// sampleAudio = nil
				// if len(matches) > 0 {
				// 	fmt.Println("FOUND A MATCH! - ", matches)
				// 	jsonData, err := json.Marshal(matches)
				// 	if err != nil {
				// 		fmt.Println(err)
				// 		return "", nil
				// 	}
				// 	return string(jsonData), nil
				// }
			}
		case <-stop:
			// Stop after 15 seconds
			fmt.Println("Stopped after 15 seconds")
			return matches, nil
		default:
			// Read RTP packets and accumulate sampleAudio
			rtpPacket, _, err := track.ReadRTP()
			if err != nil {
				if err != io.EOF {
					return nil, fmt.Errorf("error reading RTP packet: %d", err)
				}
				return nil, err
			}

			// Extract audio payload from RTP packet
			payload := rtpPacket.Payload

			sampleAudio = append(sampleAudio, payload...)
		}
	}
}

// nolint:gocognit
func SetupWebRTC(encodedOffer string) *webrtc.PeerConnection {
	// Everything below is the Pion WebRTC API! Thanks for using it ❤️.

	// Create a MediaEngine object to configure the supported codec
	m := &webrtc.MediaEngine{}

	// Setup the codecs you want to use.
	// We'll use Opus, but you can also define your own
	if err := m.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus, ClockRate: 44100, Channels: 1, SDPFmtpLine: "", RTCPFeedback: nil},
		PayloadType:        111,
	}, webrtc.RTPCodecTypeAudio); err != nil {
		panic(err)
	}

	// Create a InterceptorRegistry. This is the user configurable RTP/RTCP Pipeline.
	// This provides NACKs, RTCP Reports and other features.
	i := &interceptor.Registry{}

	// Register a intervalpli factory
	// This interceptor sends a PLI every 3 seconds. A PLI causes a keyframe to be generated by the sender.
	intervalPliFactory, err := intervalpli.NewReceiverInterceptor()
	if err != nil {
		panic(err)
	}
	i.Add(intervalPliFactory)

	// Use the default set of Interceptors
	if err = webrtc.RegisterDefaultInterceptors(m, i); err != nil {
		panic(err)
	}

	// Create the API object with the MediaEngine
	api := webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(i))

	// Prepare the configuration
	config := webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{
				URLs: []string{"stun:stun.l.google.com:19302"},
			},
		},
	}

	// Create a new RTCPeerConnection
	peerConnection, err := api.NewPeerConnection(config)
	if err != nil {
		panic(err)
	}

	// Wait for the offer to be pasted
	offer := webrtc.SessionDescription{}
	Decode(encodedOffer, &offer)

	// Set the remote SessionDescription
	err = peerConnection.SetRemoteDescription(offer)
	if err != nil {
		panic(err)
	}

	// Create answer
	answer, err := peerConnection.CreateAnswer(nil)
	if err != nil {
		panic(err)
	}

	// Create channel that is blocked until ICE Gathering is complete
	gatherComplete := webrtc.GatheringCompletePromise(peerConnection)

	// Sets the LocalDescription, and starts our UDP listeners
	err = peerConnection.SetLocalDescription(answer)
	if err != nil {
		panic(err)
	}

	// Block until ICE Gathering is complete, disabling trickle ICE
	// we do this because we only can exchange one signaling message
	// in a production application you should exchange ICE Candidates via OnICECandidate
	<-gatherComplete

	return peerConnection
}