package formats

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

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

	return h265Format, h265Media, nil
}

// ProcessH265RGBA processes H265 flow and returns PTS and IMG.
func ProcessH265RGBA(
	c *gortsplib.Client,
	h265Media *description.Media,
	h265RTPDec *rtph265.Decoder,
	h265Dec *H265Decoder,
	pkt *rtp.Packet,
	firstRandomAccess bool,
) (
	int64,
	*image.RGBA,
	error) {
	// Decode timestamp.
	pts, ok := c.PacketPTS2(h265Media, pkt)
	if !ok {
		return 0, nil, fmt.Errorf("waiting for timestamp\n")
	}

	// Extract access units from RTP packets.
	au, err := h265RTPDec.Decode(pkt)
	if err != rtph265.ErrNonStartingPacketAndNoPrevious && err != rtph265.ErrMorePacketsNeeded {
		return 0, nil, fmt.Errorf("decoding RTP packet error: %w", err)
	}

	// Wait for a random access unit.
	if !firstRandomAccess && !h265.IsRandomAccess(au) {
		return 0, nil, nil
	}
	firstRandomAccess = true

	// Convert H265 access units into RGBA frames.
	img, err := h265Dec.Decode(au)
	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
}
