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

package flexfec

import (
	
	
)

// Maximum number of media packets that can be protected by a single FEC packet.
// We are not supporting the possibility of having an FEC packet protect multiple
// SSRC source packets for now.
// https://datatracker.ietf.org/doc/html/rfc8627#section-4.2.2.1
const (
	MaxMediaPackets uint32 = 110
	MaxFecPackets   uint32 = MaxMediaPackets
)

// ProtectionCoverage defines the map of RTP packets that individual Fec packets protect.
type ProtectionCoverage struct {
	// Array of masks, each mask capable of covering up to maxMediaPkts = 110.
	// A mask is represented as a grouping of bytes where each individual bit
	// represents the coverage for the media packet at the corresponding index.
	packetMasks     [MaxFecPackets]util.BitArray
	numFecPackets   uint32
	numMediaPackets uint32
	mediaPackets    []rtp.Packet
}

// NewCoverage returns a new ProtectionCoverage object. numFecPackets represents the number of
// Fec packets that we will be generating to cover the list of mediaPackets. This allows us to know
// how big the underlying map should be.
func ( []rtp.Packet,  uint32) *ProtectionCoverage {
	 := uint32(len()) //nolint:gosec // G115

	// Basic sanity checks
	if  <= 0 ||  > MaxMediaPackets {
		return nil
	}

	// We allocate the biggest array of bitmasks that respects the max constraints.
	var  [MaxFecPackets]util.BitArray
	for  := 0;  < int(MaxFecPackets); ++ {
		[] = util.BitArray{}
	}

	 := &ProtectionCoverage{
		packetMasks:     ,
		numFecPackets:   0,
		numMediaPackets: 0,
		mediaPackets:    nil,
	}

	.UpdateCoverage(, )

	return 
}

// UpdateCoverage updates the ProtectionCoverage object with new bitmasks accounting for the numFecPackets
// we want to use to protect the batch media packets.
func ( *ProtectionCoverage) ( []rtp.Packet,  uint32) {
	 := uint32(len()) //nolint:gosec // G115

	// Basic sanity checks
	if  <= 0 ||  > MaxMediaPackets {
		return
	}

	.mediaPackets = 

	if  == .numFecPackets &&  == .numMediaPackets {
		// We have the same number of FEC packets covering the same number of media packets, we can simply
		// reuse the previous coverage map with the updated media packets.
		return
	}

	.numFecPackets = 
	.numMediaPackets = 

	// The number of FEC packets and/or the number of packets has changed, we need to update the coverage map
	// to reflect these new values.
	.resetCoverage()

	// Generate FEC bit mask where numFecPackets FEC packets are covering numMediaPackets Media packets.
	// In the packetMasks array, each FEC packet is represented by a single BitArray, each bit in a given BitArray
	// corresponds to a specific Media packet.
	// Ex: Row I, Col J is set to 1 -> FEC packet I will protect media packet J.
	for  := uint32(0);  < ; ++ {
		// We use an interleaved method to determine coverage. Given N FEC packets, Media packet X will be
		// covered by FEC packet X % N.
		 := 
		for  <  {
			.packetMasks[].SetBit()
			 += 
		}
	}
}

// ResetCoverage clears the underlying map so that we can reuse it for new batches of RTP packets.
func ( *ProtectionCoverage) () {
	for  := uint32(0);  < MaxFecPackets; ++ {
		.packetMasks[].Reset()
	}
}

// GetCoveredBy returns an iterator over RTP packets that are protected by the specified Fec packet index.
func ( *ProtectionCoverage) ( uint32) *util.MediaPacketIterator {
	 := make([]uint32, 0, .numMediaPackets)
	for  := uint32(0);  < .numMediaPackets; ++ {
		if .packetMasks[].GetBit() == 1 {
			 = append(, )
		}
	}

	return util.NewMediaPacketIterator(.mediaPackets, )
}

// ExtractMask1 returns the first section of the bitmask as defined by the FEC header.
// https://datatracker.ietf.org/doc/html/rfc8627#section-4.2.2.1
func ( *ProtectionCoverage) ( uint32) uint16 {
	return extractMask1(.packetMasks[])
}

// ExtractMask2 returns the second section of the bitmask as defined by the FEC header.
// https://datatracker.ietf.org/doc/html/rfc8627#section-4.2.2.1
func ( *ProtectionCoverage) ( uint32) uint32 {
	return extractMask2(.packetMasks[])
}

// ExtractMask3 returns the third section of the bitmask as defined by the FEC header.
// https://datatracker.ietf.org/doc/html/rfc8627#section-4.2.2.1
func ( *ProtectionCoverage) ( uint32) uint64 {
	return extractMask3(.packetMasks[])
}

// ExtractMask3_03 returns the third section of the bitmask as defined by the FEC header.
// https://datatracker.ietf.org/doc/html/draft-ietf-payload-flexible-fec-scheme-03#section-4.2
func ( *ProtectionCoverage) ( uint32) uint64 {
	return extractMask3_03(.packetMasks[])
}

func extractMask1( util.BitArray) uint16 {
	// We get the first 16 bits (64 - 16 -> shift by 48) and we shift once more for K field
	 := .Lo >> 49

	return uint16() //nolint:gosec // G115
}

func extractMask2( util.BitArray) uint32 {
	// We remove the first 15 bits
	 := .Lo << 15
	// We get the first 31 bits (64 - 32 -> shift by 32) and we shift once more for K field
	 >>= 33

	return uint32() //nolint:gosec
}

func extractMask3( util.BitArray) uint64 {
	// We remove the first 46 bits
	 := .Lo << 46
	 := .Hi >> 18
	 :=  | 

	return 
}

func extractMask3_03( util.BitArray) uint64 {
	// We remove the first 46 bits
	 := .Lo << 46
	 := .Hi >> 18
	 :=  | 
	// We shift once for the K bit.
	 >>= 1

	return 
}