rtsp_reader-writer/writer/internal/ingest/rtsp/rtsp.go

package rtsp

import (
	"errors"
	"fmt"
	"git.insit.tech/psa/rtsp_reader-writer/writer/internal/storage"
	"git.insit.tech/psa/rtsp_reader-writer/writer/pkg/converter"
	"git.insit.tech/sas/rtsp_proxy/protos/gens"
	"github.com/bluenviron/gortsplib/v4"
	"github.com/bluenviron/gortsplib/v4/pkg/base"
	"github.com/bluenviron/gortsplib/v4/pkg/description"
	"github.com/bluenviron/gortsplib/v4/pkg/format"
	"github.com/pion/rtp"
	_ "github.com/zaf/g711"
	"log"
	"os"
	"strings"
	"time"
)

// StartWriter starts the program.
func StartWriter(dir string, period int, URI string) error {
	err := RTSP(dir, period, URI)
	if err != nil {

		// Temporary solution for inner cameras.
		//
		// Change domain if a camera was flipped to another domain.
		err = changeDomain(dir, period, URI, err)
		if err != nil {
			return fmt.Errorf("change domain error: %w", err)
		}

	}
	return nil
}

// RTSP processes RTSP protocol.
func RTSP(dir string, period int, link string) error {
	resolutions := []string{"1280x720"}

	// Create file name structure and directory for files.
	cuttedURI := storage.CutURI(link)
	fn := storage.CreateFileName(dir, resolutions, cuttedURI, period)

	err := os.MkdirAll(fmt.Sprintf("%s", fn.Path), 0755)
	if err != nil {
		return fmt.Errorf("mkdirall error: %w", err)
	}

	// Create M3U8 playlist.
	go gens.MediaPlaylistGenerator(dir+"/data/"+cuttedURI, "", fn.Duration, resolutions)

	// Connect to the server.
	c := gortsplib.Client{
		UserAgent: "PSA",
	}

	u, err := base.ParseURL(link)
	if err != nil {
		return fmt.Errorf("parse URL error: %w", err)
	}

	err = c.Start(u.Scheme, u.Host)
	if err != nil {
		return fmt.Errorf("connect to the server error: %w", err)
	}
	defer c.Close()

	desc, r, err := c.Describe(u)
	if err != nil || desc == nil {
		log.Printf("medias not found for camera [%s]: %v", link, err)
		return nil
	}

	///////////////

	// Find formats.
	//var run bool
	h264Format, h264Media, err := CheckH264Format(desc)
	if err != nil {
		log.Printf("H264 format not found: %v", err)
		//run = true
	}

	//g711Format, g711Media, err := CheckG711Format(desc)
	//if err != nil {
	//	log.Printf("G711 format not found: %v", err)
	//	// run = true
	//}

	// Initialising variable for AAC.
	// var mpeg4AudioFormat *format.MPEG4Audio

	////////////////////////////////////////////////////////////////////////////////////////
	// Create RTP -> H264 decoder.
	h264RTPDec, err := h264Format.CreateDecoder()
	if err != nil {
		log.Printf("create H264 decoder error: %v", err)
	}

	// Create RTP -> H264 decoder.
	//g711RTPDec, err := g711Format.CreateDecoder()
	//if err != nil {
	//	log.Printf("create G711 decoder error: %v", err)
	//}

	////////////////////////////////////////////////////////////////////////////////////////

	// Wait for the next period.
	storage.WaitPeriod(period)
	log.Println("Start recording")

	////////////////////////////////////////////////////////////////////////////////////////

	// Setup MPEG-TS muxer.
	currentMpegtsMuxer := converter.MpegtsMuxer{
		FileName:   fn.SetNumNTime(),
		H264Format: h264Format,
		// Mpeg4AudioFormat: mpeg4AudioFormat,
	}

	err = currentMpegtsMuxer.Initialize()
	if err != nil {
		panic(err)
	}
	defer currentMpegtsMuxer.Close()

	////////////////////////////////////////////////////////////////////////////////////////

	// Setup all medias.
	err = c.SetupAll(desc.BaseURL, desc.Medias)
	if err != nil {
		panic(err)
	}

	////////////////////////////////////////////////////////////////////////////////////////

	// Called when a H264/RTP or G711/RTP packet arrives.
	c.OnPacketRTPAny(func(medi *description.Media, forma format.Format, pkt *rtp.Packet) {
		switch forma.(type) {
		case *format.H264:
			// Process H264 flow and return PTS and AU.
			pts, au, err := ProcH264(&c, h264Media, h264RTPDec, pkt)
			if err != nil {
				//log.Printf("%s: process H264 error: %s\n", cuttedURI, err)
			}

			// Encode the access unit into MPEG-TS.
			err = currentMpegtsMuxer.WriteH264(au, pts)
			if err != nil {
				//log.Printf("%s: writing H264 packet: %s\n", cuttedURI, err)
			}

		case *format.G711:
			//// Process G711 flow and returns PTS and AU.
			//_, _, err := media.ProcG711(&c, g711Media, g711RTPDec, pkt)
			//if err != nil {
			//	log.Printf("%s: process G711 error: %s\n", cuttedURI, err)
			//}

			//// Convert G711 to AAC.
			//_, err = converter.ConvertG711ToAAC(au, f.MULaw) // take aacAu
			//if err != nil {
			//	log.Printf("%s: converting G711 to AAC frame: %s\n", cuttedURI, err)
			//}

			/*
				// Encode the access unit into MPEG-TS.
				err = MpegtsMuxer.writeMPEG4Audio([][]byte{aacAu}, pts)
				if err != nil {
					log.Printf("MPEG-TS write error: %v", err)
					return
				}
			*/
		}
	})

	////////////////////////////////////////////////////////////////////////////////////////

	// Send PLAY request.
	_, err = c.Play(nil)
	if err != nil {
		log.Fatalln("Ошибка запуска воспроизведения:", err)
	}

	// Create ticker for rotation files.
	ticker := time.NewTicker(time.Duration(period) * time.Second)
	defer ticker.Stop()

	// Rotate files.
	go func() {
		for range ticker.C {
			currentMpegtsMuxer.Close()
			currentMpegtsMuxer.FileName = fn.SetNumNTime()

			err = currentMpegtsMuxer.Initialize()
			if err != nil {
				panic(err)
			}

			log.Println("New file for recording created:", currentMpegtsMuxer.FileName)
		}
	}()

	// Restart program if client gets TEARDOWN request.
	c.OnRequest = func(req *base.Request) {
		if req.Method == base.Teardown {
			log.Printf("got TEARDOWN request from server: %v", req)

			//err = procRTSP(period, URI)
			//if err != nil {
			//	log.Fatalf("restart RTSP error: %s\n", err)
			//}
		}
	}
	panic(c.Wait())
}

// changeDomain changes domain if a camera was flipped to another domain.
func changeDomain(directory string, period int, URI string, err error) error {
	err2 := errors.New("404 (Not found)")
	if errors.As(err, &err2) {
		if strings.Contains(URI, "video-1") {
			err = RTSP(directory, period, strings.Replace(URI, "video-1", "video-2", 1))
			if err != nil {
				return err
			}
		} else {
			err = RTSP(directory, period, strings.Replace(URI, "video-2", "video-1", 1))
			if err != nil {
				return err
			}
		}
	} else {
		panic(err)
	}

	return nil
}