package formats

import (
	"C"
	"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/rtpav1"
	"github.com/bluenviron/mediacommon/v2/pkg/codecs/av1"
	"github.com/pion/rtp"
	"image"
)

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

	return av1Format, av1Media, nil
}

// ProcessAV1 processes AV1 flow and returns PTS and IMG.
func ProcessAV1(
	c *gortsplib.Client,
	av1Media *description.Media,
	av1RTPDec *rtpav1.Decoder,
	pkt *rtp.Packet,
	av1Dec *AV1Decoder,
	firstRandomReceived bool,
) (
	int64,
	*image.RGBA,
	error) {
	// Decode timestamp.
	pts, ok := c.PacketPTS2(av1Media, pkt)
	if !ok {
		return 0, nil, fmt.Errorf("waiting for timestamp\n")
	}

	// Extract AV1 temporal units from RTP packets.
	tu, err := av1RTPDec.Decode(pkt)
	if err != rtpav1.ErrNonStartingPacketAndNoPrevious && err != rtpav1.ErrMorePacketsNeeded {
		return 0, nil, fmt.Errorf("decoding RTP packet error: %w", err)
	}

	// Wait for a random access unit.
	IsRandomAccess2, _ := av1.IsRandomAccess(tu)
	if !firstRandomReceived && !IsRandomAccess2 {
		return 0, nil, fmt.Errorf("waiting for a random access unit\n")
	}
	firstRandomReceived = true

	// Convert AV1 temporal units into RGBA frames.
	img, err := av1Dec.Decode(tu)
	if err != nil {
		return 0, nil, fmt.Errorf("convert into RGBA frames error\n")
	}

	// Wait for a frame.
	if img == nil {
		return 0, nil, fmt.Errorf("frame not found\n")
	}

	return pts, img, nil
}