package formats

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/bluenviron/gortsplib/v4"
	"github.com/bluenviron/gortsplib/v4/pkg/description"
	"github.com/bluenviron/gortsplib/v4/pkg/format"
	"github.com/bluenviron/gortsplib/v4/pkg/format/rtpmjpeg"
	"github.com/pion/rtp"
	"image"
	"image/jpeg"
)

// FindMJPEGFormat finds the MJPEG media and format.
func FindMJPEGFormat(desc *description.Session) (*format.MJPEG, *description.Media, error) {
	var mjpegFormat *format.MJPEG
	mjpegMedia := desc.FindFormat(&mjpegFormat)
	if mjpegMedia == nil {
		return nil, nil, errors.New("media MJPEG not found")
	}

	return mjpegFormat, mjpegMedia, nil
}

// ProcessMJPEGRGBA processes MJPEG flow and returns PTS and IMG.
func ProcessMJPEGRGBA(
	c *gortsplib.Client,
	mjpegMedia *description.Media,
	mjpegRTPDec *rtpmjpeg.Decoder,
	pkt *rtp.Packet,
	t string,
) (
	int64,
	image.Image,
	error) {
	// Decode timestamp.
	pts, ok := c.PacketPTS2(mjpegMedia, pkt)
	if !ok {
		return 0, nil, fmt.Errorf("[%v]: waiting for timestamp\n", t)
	}

	// Extract JPEG images from RTP packets.
	enc, err := mjpegRTPDec.Decode(pkt)
	if err != rtpmjpeg.ErrNonStartingPacketAndNoPrevious && err != rtpmjpeg.ErrMorePacketsNeeded {
		return 0, nil, fmt.Errorf("[%v]: decoding RTP packet error: %w", t, err)
	}

	// Convert JPEG images into RGBA frames.
	img, err := jpeg.Decode(bytes.NewReader(enc))
	if err != nil {
		return 0, nil, fmt.Errorf("[%v]: convert into RGBA frames error\n", t)
	}

	return pts, img, nil
}