// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
// SPDX-License-Identifier: MIT

package flexfec

import (
	
	

	
	
)

// streamState holds the state for a single stream.
type streamState struct {
	mu             sync.Mutex
	flexFecEncoder FlexEncoder
	packetBuffer   []rtp.Packet
}

// FecInterceptor implements FlexFec.
type FecInterceptor struct {
	interceptor.NoOp
	mu              sync.Mutex
	streams         map[uint32]*streamState
	numMediaPackets uint32
	numFecPackets   uint32
	encoderFactory  EncoderFactory
}

// FecInterceptorFactory creates new FecInterceptors.
type FecInterceptorFactory struct {
	opts []FecOption
}

// NewFecInterceptor returns a new Fec interceptor factory.
func ( ...FecOption) (*FecInterceptorFactory, error) {
	return &FecInterceptorFactory{opts: }, nil
}

// NewInterceptor constructs a new FecInterceptor.
func ( *FecInterceptorFactory) ( string) (interceptor.Interceptor, error) {
	 := &FecInterceptor{
		streams:         make(map[uint32]*streamState),
		numMediaPackets: 5,
		numFecPackets:   2,
		encoderFactory:  FlexEncoder03Factory{},
	}

	for ,  := range .opts {
		if  := ();  != nil {
			return nil, 
		}
	}

	return , nil
}

// UnbindLocalStream removes the stream state for a specific SSRC.
func ( *FecInterceptor) ( *interceptor.StreamInfo) {
	.mu.Lock()
	defer .mu.Unlock()

	delete(.streams, .SSRC)
}

// BindLocalStream lets you modify any outgoing RTP packets. It is called once for per LocalStream. The returned method
// will be called once per rtp packet.
func ( *FecInterceptor) (
	 *interceptor.StreamInfo,  interceptor.RTPWriter,
) interceptor.RTPWriter {
	if .PayloadTypeForwardErrorCorrection == 0 || .SSRCForwardErrorCorrection == 0 {
		return 
	}

	 := .SSRC

	.mu.Lock()
	 := &streamState{
		// Chromium supports version flexfec-03 of existing draft, this is the one we will configure by default
		// although we should support configuring the latest (flexfec-20) as well.
		flexFecEncoder: .encoderFactory.NewEncoder(.PayloadTypeForwardErrorCorrection, .SSRCForwardErrorCorrection),
		packetBuffer:   make([]rtp.Packet, 0),
	}
	.streams[] = 
	.mu.Unlock()

	return interceptor.RTPWriterFunc(
		func( *rtp.Header,  []byte,  interceptor.Attributes) (int, error) {
			// Ignore non-media packets
			if .SSRC !=  {
				return .Write(, , )
			}

			var  []rtp.Packet
			.mu.Lock()
			.packetBuffer = append(.packetBuffer, rtp.Packet{
				Header:  *,
				Payload: ,
			})

			// Check if we have enough packets to generate FEC
			if len(.packetBuffer) == int(.numMediaPackets) {
				 = .flexFecEncoder.EncodeFec(.packetBuffer, .numFecPackets)
				// Reset the packet buffer now that we've sent the corresponding FEC packets.
				.packetBuffer = nil
			}
			.mu.Unlock()

			var  []error
			,  := .Write(, , )
			if  != nil {
				 = append(, )
			}

			for ,  := range  {
				 := .Header

				_,  = .Write(&, .Payload, )
				if  != nil {
					 = append(, )
				}
			}

			return , errors.Join(...)
		},
	)
}