package main

import (
	"errors"
	"github.com/aler9/gortsplib"
	"github.com/aler9/gortsplib/pkg/url"
	"image"
	"image/jpeg"
	"log"
	"os"
	"strconv"
	"time"
)

// This example shows how to
// 1. connect to a RTSP server and read all tracks on a path
// 2. check if there's a H264 track
// 3. decode H264 into RGBA frames
// 4. encode the frames into JPEG images and save them on disk

// This example requires the ffmpeg libraries, that can be installed in this way:
// apt install -y libavformat-dev libswscale-dev gcc pkg-config

func playRtsp(c *gortsplib.Client, u *url.URL, saveImage func(image.Image)) error {
	if tracks, baseUrl, _, e := c.Describe(u); nil != e {
		return e
	} else {
		h264TrackID, h264track := func() (int, *gortsplib.TrackH264) {
			for i, track := range tracks {
				if h264track, ok := track.(*gortsplib.TrackH264); ok {
					return i, h264track
				}
			}
			return -1, nil
		}()
		if h264TrackID < 0 {
			return errors.New("h264 track not found")
		}
		// setup H264->raw frames decoder
		h264dec, err := newH264Decoder()
		if err != nil {
			panic(err)
		}
		defer h264dec.close()

		// if present, send SPS and PPS from the SDP to the decoder
		sps := h264track.SafeSPS()
		if sps != nil {
			h264dec.decode(sps)
		}
		pps := h264track.SafePPS()
		if pps != nil {
			h264dec.decode(pps)
		}

		// called when a RTP packet arrives
		if nil != saveImage {
			log.Println("setupRtsp:> enabled saveImage")
			c.OnPacketRTP = func(ctx *gortsplib.ClientOnPacketRTPCtx) {

				if ctx.TrackID != h264TrackID {
					//log.Println("x")
					return
				}

				if ctx.H264NALUs == nil {
					//log.Println("y")
					return
				}

				for _, nalu := range ctx.H264NALUs {
					// convert H264 NALUs to RGBA frames
					img, err := h264dec.decode(nalu)
					if err != nil {
						panic(err)
					}

					// wait for a frame
					if img == nil {
						//log.Println("z")
						continue
					} else {
						//log.Println("...")
						saveImage(img)
					}
					// convert frame to JPEG and save to file

				}
			}
		}

		// setup and read all tracks
		err = c.SetupAndPlay(tracks, baseUrl)
		if err != nil {
			return err
		}
		// wait until a fatal error
		return c.Wait()
	}
}

func setupRtsp(u string) (*gortsplib.Client, *url.URL, error) {
	var c = &gortsplib.Client{}
	if u, e := url.Parse(u); nil != e {
		return nil, nil, e
	} else if e := c.Start(u.Scheme, u.Host); nil != e {
		return nil, nil, e
	} else {
		return c, u, nil
	}
}

func saveToFile(img image.Image) error {
	// create file
	fname := strconv.FormatInt(time.Now().UnixNano()/int64(time.Millisecond), 10) + ".jpg"
	f, err := os.Create(fname)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	log.Println("saving", fname)

	// convert to jpeg
	return jpeg.Encode(f, img, &jpeg.Options{
		Quality: 60,
	})
}