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

package rtp

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
)

const (
	headerExtensionIDReserved = 0xF
)

// HeaderExtension represents an RTP extension header.
type HeaderExtension interface {
	Set(id uint8, payload []byte) error
	GetIDs() []uint8
	Get(id uint8) []byte
	Del(id uint8) error

	Unmarshal(buf []byte) (int, error)
	Marshal() ([]byte, error)
	MarshalTo(buf []byte) (int, error)
	MarshalSize() int
}

// OneByteHeaderExtension is an RFC8285 one-byte header extension.
type OneByteHeaderExtension struct {
	payload []byte
}

// Set sets the extension payload for the specified ID.
func (e *OneByteHeaderExtension) Set(id uint8, buf []byte) error {
	if err := headerExtensionCheck(ExtensionProfileOneByte, id, buf); err != nil {
		return err
	}

	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n] >> 4
		payloadLen := int(e.payload[n]&^0xF0 + 1)
		n++

		if extid == id {
			e.payload = append(e.payload[:n], append(buf, e.payload[n+payloadLen:]...)...)

			return nil
		}
		n += payloadLen
	}

	if len(e.payload) == 0 && !bytes.HasPrefix(buf, []byte{0xBE, 0xDE, 0x00, 0x00}) {
		e.payload = []byte{0xBE, 0xDE, 0x00, 0x00}
	}

	e.payload = append(e.payload, (id<<4 | uint8(len(buf)-1))) // nolint: gosec // G115
	e.payload = append(e.payload, buf...)
	binary.BigEndian.PutUint16(e.payload[2:4], binary.BigEndian.Uint16(e.payload[2:4])+1)

	return nil
}

// GetIDs returns the available IDs.
func (e *OneByteHeaderExtension) GetIDs() []uint8 {
	ids := make([]uint8, 0, binary.BigEndian.Uint16(e.payload[2:4]))
	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n] >> 4
		payloadLen := int(e.payload[n]&^0xF0 + 1)
		n++

		if extid == headerExtensionIDReserved {
			break
		}

		ids = append(ids, extid)
		n += payloadLen
	}

	return ids
}

// Get returns the payload of the extension with the given ID.
func (e *OneByteHeaderExtension) Get(id uint8) []byte {
	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n] >> 4
		payloadLen := int(e.payload[n]&^0xF0 + 1)
		n++

		if extid == id {
			return e.payload[n : n+payloadLen]
		}
		n += payloadLen
	}

	return nil
}

// Del deletes the extension with the specified ID.
func (e *OneByteHeaderExtension) Del(id uint8) error {
	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n] >> 4
		payloadLen := int(e.payload[n]&^0xF0 + 1)

		if extid == id {
			e.payload = append(e.payload[:n], e.payload[n+1+payloadLen:]...)

			return nil
		}
		n += payloadLen + 1
	}

	return errHeaderExtensionNotFound
}

// Unmarshal parses the extension payload.
func (e *OneByteHeaderExtension) Unmarshal(buf []byte) (int, error) {
	profile := binary.BigEndian.Uint16(buf[0:2])
	if profile != ExtensionProfileOneByte {
		return 0, fmt.Errorf("%w actual(%x)", errHeaderExtensionNotFound, buf[0:2])
	}
	e.payload = buf

	return len(buf), nil
}

// Marshal returns the extension payload.
func (e OneByteHeaderExtension) Marshal() ([]byte, error) {
	return e.payload, nil
}

// MarshalTo writes the extension payload to the given buffer.
func (e OneByteHeaderExtension) MarshalTo(buf []byte) (int, error) {
	size := e.MarshalSize()
	if size > len(buf) {
		return 0, io.ErrShortBuffer
	}

	return copy(buf, e.payload), nil
}

// MarshalSize returns the size of the extension payload.
func (e OneByteHeaderExtension) MarshalSize() int {
	return len(e.payload)
}

// TwoByteHeaderExtension is an RFC8285 two-byte header extension.
type TwoByteHeaderExtension struct {
	payload []byte
}

// Set sets the extension payload for the specified ID.
func (e *TwoByteHeaderExtension) Set(id uint8, buf []byte) error {
	if err := headerExtensionCheck(ExtensionProfileTwoByte, id, buf); err != nil {
		return err
	}

	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n]
		n++

		payloadLen := int(e.payload[n])
		n++

		if extid == id {
			e.payload = append(e.payload[:n], append(buf, e.payload[n+payloadLen:]...)...)

			return nil
		}
		n += payloadLen
	}

	if len(e.payload) == 0 && !bytes.HasPrefix(buf, []byte{0x10, 0x00, 0x00, 0x00}) {
		e.payload = []byte{0x10, 0x00, 0x00, 0x00}
	}

	e.payload = append(e.payload, id, uint8(len(buf))) // nolint: gosec // G115
	e.payload = append(e.payload, buf...)
	binary.BigEndian.PutUint16(e.payload[2:4], binary.BigEndian.Uint16(e.payload[2:4])+1)

	return nil
}

// GetIDs returns the available IDs.
func (e *TwoByteHeaderExtension) GetIDs() []uint8 {
	ids := make([]uint8, 0, binary.BigEndian.Uint16(e.payload[2:4]))
	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n]
		n++

		payloadLen := int(e.payload[n])
		n++

		ids = append(ids, extid)
		n += payloadLen
	}

	return ids
}

// Get returns the payload of the extension with the given ID.
func (e *TwoByteHeaderExtension) Get(id uint8) []byte {
	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n]
		n++

		payloadLen := int(e.payload[n])
		n++

		if extid == id {
			return e.payload[n : n+payloadLen]
		}
		n += payloadLen
	}

	return nil
}

// Del deletes the extension with the specified ID.
func (e *TwoByteHeaderExtension) Del(id uint8) error {
	for n := 4; n < len(e.payload); {
		if e.payload[n] == 0x00 { // padding
			n++

			continue
		}

		extid := e.payload[n]

		payloadLen := int(e.payload[n+1])

		if extid == id {
			e.payload = append(e.payload[:n], e.payload[n+2+payloadLen:]...)

			return nil
		}
		n += payloadLen + 2
	}

	return errHeaderExtensionNotFound
}

// Unmarshal parses the extension payload.
func (e *TwoByteHeaderExtension) Unmarshal(buf []byte) (int, error) {
	profile := binary.BigEndian.Uint16(buf[0:2])
	if profile != ExtensionProfileTwoByte {
		return 0, fmt.Errorf("%w actual(%x)", errHeaderExtensionNotFound, buf[0:2])
	}
	e.payload = buf

	return len(buf), nil
}

// Marshal returns the extension payload.
func (e TwoByteHeaderExtension) Marshal() ([]byte, error) {
	return e.payload, nil
}

// MarshalTo marshals the extension to the given buffer.
func (e TwoByteHeaderExtension) MarshalTo(buf []byte) (int, error) {
	size := e.MarshalSize()
	if size > len(buf) {
		return 0, io.ErrShortBuffer
	}

	return copy(buf, e.payload), nil
}

// MarshalSize returns the size of the extension payload.
func (e TwoByteHeaderExtension) MarshalSize() int {
	return len(e.payload)
}

// RawExtension represents an RFC3550 header extension.
type RawExtension struct {
	payload []byte
}

// Set sets the extension payload for the specified ID.
func (e *RawExtension) Set(id uint8, payload []byte) error {
	if err := headerExtensionCheck(0, id, payload); err != nil {
		return err
	}

	e.payload = payload

	return nil
}

// GetIDs returns the available IDs.
func (e *RawExtension) GetIDs() []uint8 {
	return []uint8{0}
}

// Get returns the payload of the extension with the given ID.
func (e *RawExtension) Get(id uint8) []byte {
	if id == 0 {
		return e.payload
	}

	return nil
}

// Del deletes the extension with the specified ID.
func (e *RawExtension) Del(id uint8) error {
	if id == 0 {
		e.payload = nil

		return nil
	}

	return fmt.Errorf("%w actual(%d)", errRFC3550HeaderIDRange, id)
}

// Unmarshal parses the extension from the given buffer.
func (e *RawExtension) Unmarshal(buf []byte) (int, error) {
	profile := binary.BigEndian.Uint16(buf[0:2])
	if profile == ExtensionProfileOneByte || profile == ExtensionProfileTwoByte {
		return 0, fmt.Errorf("%w actual(%x)", errHeaderExtensionNotFound, buf[0:2])
	}
	e.payload = buf

	return len(buf), nil
}

// Marshal returns the raw extension payload.
func (e RawExtension) Marshal() ([]byte, error) {
	return e.payload, nil
}

// MarshalTo marshals the extension to the given buffer.
func (e RawExtension) MarshalTo(buf []byte) (int, error) {
	size := e.MarshalSize()
	if size > len(buf) {
		return 0, io.ErrShortBuffer
	}

	return copy(buf, e.payload), nil
}

// MarshalSize returns the size of the extension when marshaled.
func (e RawExtension) MarshalSize() int {
	return len(e.payload)
}

// Assert that id + value is valid for give Header Extension Profile.
func headerExtensionCheck(extensionProfile uint16, id uint8, payload []byte) error {
	switch extensionProfile {
	// RFC 8285 RTP One Byte Header Extension
	case ExtensionProfileOneByte:
		if id < 1 || id > 14 {
			return fmt.Errorf("%w actual(%d)", errRFC8285OneByteHeaderIDRange, id)
		}
		if len(payload) > 16 {
			return fmt.Errorf("%w actual(%d)", errRFC8285OneByteHeaderSize, len(payload))
		}
	// RFC 8285 RTP Two Byte Header Extension
	case ExtensionProfileTwoByte:
		if id < 1 {
			return fmt.Errorf("%w actual(%d)", errRFC8285TwoByteHeaderIDRange, id)
		}
		if len(payload) > 255 {
			return fmt.Errorf("%w actual(%d)", errRFC8285TwoByteHeaderSize, len(payload))
		}
	default: // RFC3550 Extension
		if id != 0 {
			return fmt.Errorf("%w actual(%d)", errRFC3550HeaderIDRange, id)
		}
	}

	return nil
}
