package gateway

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"errors"
)

const (
	// Apollo limits clear audio payloads to 1400 bytes and uses AES-CBC with
	// PKCS#7 padding. FEC adds its fixed 12-byte header outside that ciphertext.
	// This is the largest accepted UDP datagram, not an allocation hint.
	apolloMediaMaximumPacket = 12 + 12 + 1408
	apolloVideoHeaderSize    = 32
	apolloRTPHeaderSize      = 12
	apolloVideoNVHeaderSize  = 16
	apolloVideoRawPacketSize = 1024 + 16
)

var (
	errApolloMedia       = errors.New("apollo media malformed")
	errApolloMediaParity = errors.New("apollo media parity packet")
)

type apolloMediaCodec struct {
	block cipher.Block
	aead  cipher.AEAD
	keyID uint32
}

type apolloRTPPacket struct {
	extension   bool
	payloadType byte
	sequence    uint16
	payload     []byte
}

type apolloAudioShard struct {
	sequence    uint16
	timestamp   uint32
	ssrc        uint32
	base        uint16
	parityIndex uint8
	parity      bool
	payload     []byte
}

func newApolloMediaCodec(key []byte, keyID uint32) (*apolloMediaCodec, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return &apolloMediaCodec{block: block, aead: aead, keyID: keyID}, nil
}

func apolloMediaPing(payload []byte, sequence uint32) []byte {
	if len(payload) != 16 {
		return nil
	}
	ping := make([]byte, 20)
	copy(ping, payload)
	binary.BigEndian.PutUint32(ping[16:], sequence)
	return ping
}

func (c *apolloMediaCodec) OpenVideo(packet []byte) (apolloVideoShard, error) {
	if c == nil || c.aead == nil || len(packet) != apolloVideoHeaderSize+apolloVideoRawPacketSize {
		return apolloVideoShard{}, errApolloMedia
	}
	sealed := make([]byte, len(packet)-apolloVideoHeaderSize+apolloControlTagSize)
	copy(sealed, packet[apolloVideoHeaderSize:])
	copy(sealed[len(packet)-apolloVideoHeaderSize:], packet[16:apolloVideoHeaderSize])
	plaintext, err := c.aead.Open(nil, packet[:12], sealed, nil)
	if err != nil {
		return apolloVideoShard{}, errApolloMedia
	}
	rtp, err := parseApolloRTP(plaintext)
	if err != nil || !rtp.extension || rtp.payloadType != 0 || len(rtp.payload) != 1024 {
		return apolloVideoShard{}, errApolloMedia
	}
	nv := rtp.payload[:apolloVideoNVHeaderSize]
	flags := nv[8]
	fecInfo := binary.LittleEndian.Uint32(nv[12:16])
	dataPackets := int(fecInfo >> 22)
	fecIndex := int((fecInfo >> 12) & 0x03ff)
	fecPercent := int((fecInfo >> 4) & 0xff)
	if dataPackets < 1 || dataPackets > apolloVideoMaximumDataShards {
		return apolloVideoShard{}, errApolloMedia
	}
	parityPackets := (dataPackets*fecPercent + 99) / 100
	if dataPackets+parityPackets > 255 || fecIndex >= dataPackets+parityPackets {
		return apolloVideoShard{}, errApolloMedia
	}
	block := (nv[11] >> 4) & 0x03
	lastBlock := (nv[11] >> 6) & 0x03
	if block > lastBlock {
		return apolloVideoShard{}, errApolloMedia
	}
	return apolloVideoShard{
		frame:       binary.LittleEndian.Uint32(nv[4:8]),
		block:       block,
		lastBlock:   lastBlock,
		dataPackets: dataPackets,
		parity:      parityPackets,
		index:       fecIndex,
		sequence:    rtp.sequence,
		streamIndex: binary.LittleEndian.Uint32(nv[:4]) >> 8,
		flags:       flags,
		payload:     append([]byte(nil), rtp.payload[apolloVideoNVHeaderSize:]...),
	}, nil
}

func (c *apolloMediaCodec) OpenAudio(packet []byte) (apolloAudioShard, error) {
	if c == nil || c.block == nil || len(packet) <= apolloRTPHeaderSize || len(packet) > apolloMediaMaximumPacket {
		return apolloAudioShard{}, errApolloMedia
	}
	rtp, err := parseApolloRTP(packet)
	if err != nil || rtp.extension || len(rtp.payload) == 0 {
		return apolloAudioShard{}, errApolloMedia
	}
	if rtp.payloadType == 97 && len(rtp.payload)%aes.BlockSize == 0 {
		return apolloAudioShard{
			sequence: rtp.sequence, timestamp: binary.BigEndian.Uint32(packet[4:8]), ssrc: binary.BigEndian.Uint32(packet[8:12]),
			base: rtp.sequence &^ 3, payload: append([]byte(nil), rtp.payload...),
		}, nil
	}
	if rtp.payloadType != 127 || len(rtp.payload) <= 12 || len(rtp.payload)-12 > 1408 || rtp.payload[0] > 1 || rtp.payload[1] != 97 {
		return apolloAudioShard{}, errApolloMedia
	}
	base := binary.BigEndian.Uint16(rtp.payload[2:4])
	if base&3 != 0 || len(rtp.payload[12:])%aes.BlockSize != 0 {
		return apolloAudioShard{}, errApolloMedia
	}
	return apolloAudioShard{
		sequence: rtp.sequence, timestamp: binary.BigEndian.Uint32(rtp.payload[4:8]), ssrc: binary.BigEndian.Uint32(rtp.payload[8:12]),
		base: base, parityIndex: rtp.payload[0], parity: true, payload: append([]byte(nil), rtp.payload[12:]...),
	}, nil
}

func (c *apolloMediaCodec) openApolloAudioCipher(sequence uint16, payload []byte) ([]byte, error) {
	if c == nil || c.block == nil || len(payload) == 0 || len(payload) > 1408 || len(payload)%aes.BlockSize != 0 {
		return nil, errApolloMedia
	}
	plaintext := append([]byte(nil), payload...)
	iv := make([]byte, aes.BlockSize)
	binary.BigEndian.PutUint32(iv, c.keyID+uint32(sequence))
	cipher.NewCBCDecrypter(c.block, iv).CryptBlocks(plaintext, plaintext)
	padding := int(plaintext[len(plaintext)-1])
	if padding == 0 || padding > aes.BlockSize || padding > len(plaintext) {
		return nil, errApolloMedia
	}
	for _, value := range plaintext[len(plaintext)-padding:] {
		if int(value) != padding {
			return nil, errApolloMedia
		}
	}
	return append([]byte(nil), plaintext[:len(plaintext)-padding]...), nil
}

func parseApolloRTP(packet []byte) (apolloRTPPacket, error) {
	if len(packet) < apolloRTPHeaderSize || packet[0]>>6 != 2 || packet[0]&0x2f != 0 {
		return apolloRTPPacket{}, errApolloMedia
	}
	offset := apolloRTPHeaderSize
	extension := packet[0]&0x10 != 0
	if extension {
		if len(packet) < offset+4 || binary.BigEndian.Uint16(packet[14:16]) != 0 {
			return apolloRTPPacket{}, errApolloMedia
		}
		offset += 4
	}
	if offset >= len(packet) {
		return apolloRTPPacket{}, errApolloMedia
	}
	return apolloRTPPacket{extension: extension, payloadType: packet[1] & 0x7f, sequence: binary.BigEndian.Uint16(packet[2:4]), payload: packet[offset:]}, nil
}