package rtpproc

import (
	"fmt"
	"time"

	"github.com/pion/rtp"

	"github.com/aler9/gortsplib/pkg/h264"
	"github.com/aler9/gortsplib/pkg/rtph264"
)

const (
	// 1500 (UDP MTU) - 20 (IP header) - 8 (UDP header)
	maxPacketSize = 1472
)

// ProcessorOutput is the output of Process().
type ProcessorOutput struct {
	Packet       *rtp.Packet
	PTSEqualsDTS bool
	H264NALUs    [][]byte
	H264PTS      time.Duration
}

// Processor is used to process incoming RTP packets, in order to:
// - remove padding
// - decode packets encoded with supported codecs
// - re-encode packets if they are bigger than maximum allowed.
type Processor struct {
	isH264 bool
	isTCP  bool

	h264Decoder *rtph264.Decoder
	h264Encoder *rtph264.Encoder
}

// NewProcessor allocates a Processor.
func NewProcessor(isH264 bool, isTCP bool) *Processor {
	p := &Processor{
		isH264: isH264,
		isTCP:  isTCP,
	}

	if isH264 {
		p.h264Decoder = &rtph264.Decoder{}
		p.h264Decoder.Init()
	}

	return p
}

func (p *Processor) processH264(pkt *rtp.Packet) ([]*ProcessorOutput, error) {
	// decode
	nalus, pts, err := p.h264Decoder.DecodeUntilMarker(pkt)
	if err != nil {
		if err == rtph264.ErrNonStartingPacketAndNoPrevious ||
			err == rtph264.ErrMorePacketsNeeded {
			return []*ProcessorOutput{{
				Packet:       pkt,
				PTSEqualsDTS: false,
			}}, nil
		}
		return nil, err
	}
	ptsEqualsDTS := h264.IDRPresent(nalus)

	// re-encode if packets use non-standard sizes
	if p.isTCP && p.h264Encoder == nil && pkt.MarshalSize() > maxPacketSize {
		v1 := pkt.SSRC
		v2 := pkt.SequenceNumber
		v3 := pkt.Timestamp
		p.h264Encoder = &rtph264.Encoder{
			PayloadType:           pkt.PayloadType,
			SSRC:                  &v1,
			InitialSequenceNumber: &v2,
			InitialTimestamp:      &v3,
		}
		p.h264Encoder.Init()
	}

	if p.h264Encoder != nil {
		packets, err := p.h264Encoder.Encode(nalus, pts)
		if err != nil {
			return nil, err
		}

		output := make([]*ProcessorOutput, len(packets))

		for i, pkt := range packets {
			if i != len(packets)-1 {
				output[i] = &ProcessorOutput{
					Packet:       pkt,
					PTSEqualsDTS: false,
				}
			} else {
				output[i] = &ProcessorOutput{
					Packet:       pkt,
					PTSEqualsDTS: ptsEqualsDTS,
				}
			}
		}

		return output, nil
	}

	return []*ProcessorOutput{{
		Packet:       pkt,
		PTSEqualsDTS: ptsEqualsDTS,
		H264NALUs:    nalus,
		H264PTS:      pts,
	}}, nil
}

// Process processes a RTP packet.
func (p *Processor) Process(pkt *rtp.Packet) ([]*ProcessorOutput, error) {
	// remove padding
	pkt.Header.Padding = false
	pkt.PaddingSize = 0

	if p.h264Decoder != nil {
		return p.processH264(pkt)
	}

	if p.isTCP && pkt.MarshalSize() > maxPacketSize {
		return nil, fmt.Errorf("payload size (%d) greater than maximum allowed (%d)",
			pkt.MarshalSize(), maxPacketSize)
	}

	return []*ProcessorOutput{{
		Packet:       pkt,
		PTSEqualsDTS: true,
	}}, nil
}