package gateway

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

const (
	apolloControlOuterType    = 0x0001
	apolloControlHeaderSize   = 8
	apolloControlTagSize      = 16
	apolloControlInnerSize    = 4
	apolloControlMaximumPlain = 2048

	apolloControlTypeStart  = 0x0307
	apolloControlTypeIDR    = 0x0302
	apolloControlTypePing   = 0x0200
	apolloControlTypeInput  = 0x0206
	apolloControlTypeFEC    = 0x5502
	apolloControlTypeRumble = 0x010b
	apolloControlTypeHDR    = 0x010e
	apolloControlTypeTerm   = 0x0109
)

var errApolloControl = errors.New("apollo control malformed")

type apolloControlMessage struct {
	typeID  uint16
	payload []byte
}

type apolloControlCodec struct {
	aead       cipher.AEAD
	nextClient uint32
	lastHost   uint32
	hostSeen   bool
}

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

func (c *apolloControlCodec) SealClient(typeID uint16, payload []byte) ([]byte, error) {
	if c == nil || c.aead == nil || typeID == 0 || len(payload) > apolloControlMaximumPlain || c.nextClient == ^uint32(0) {
		return nil, errApolloControl
	}
	sequence := c.nextClient
	c.nextClient++
	inner := make([]byte, apolloControlInnerSize+len(payload))
	binary.LittleEndian.PutUint16(inner[:2], typeID)
	binary.LittleEndian.PutUint16(inner[2:4], uint16(len(payload)))
	copy(inner[4:], payload)
	nonce := apolloControlNonce(sequence, 'C')
	sealed := c.aead.Seal(nil, nonce[:], inner, nil)
	packet := make([]byte, apolloControlHeaderSize+len(sealed))
	binary.LittleEndian.PutUint16(packet[:2], apolloControlOuterType)
	binary.LittleEndian.PutUint16(packet[2:4], uint16(4+len(sealed)))
	binary.LittleEndian.PutUint32(packet[4:8], sequence)
	copy(packet[8:24], sealed[len(inner):])
	copy(packet[24:], sealed[:len(inner)])
	return packet, nil
}

func (c *apolloControlCodec) OpenHost(packet []byte) (apolloControlMessage, error) {
	if c == nil || c.aead == nil || len(packet) < apolloControlHeaderSize+apolloControlTagSize+apolloControlInnerSize || len(packet) > apolloControlHeaderSize+apolloControlTagSize+apolloControlMaximumPlain {
		return apolloControlMessage{}, errApolloControl
	}
	if binary.LittleEndian.Uint16(packet[:2]) != apolloControlOuterType || int(binary.LittleEndian.Uint16(packet[2:4])) != len(packet)-4 {
		return apolloControlMessage{}, errApolloControl
	}
	sequence := binary.LittleEndian.Uint32(packet[4:8])
	if c.hostSeen && sequence <= c.lastHost {
		return apolloControlMessage{}, errApolloControl
	}
	nonce := apolloControlNonce(sequence, 'H')
	sealed := make([]byte, len(packet)-apolloControlHeaderSize)
	copy(sealed, packet[24:])
	copy(sealed[len(packet)-24:], packet[8:24])
	plaintext, err := c.aead.Open(nil, nonce[:], sealed, nil)
	if err != nil || len(plaintext) < apolloControlInnerSize {
		return apolloControlMessage{}, errApolloControl
	}
	length := int(binary.LittleEndian.Uint16(plaintext[2:4]))
	if length != len(plaintext)-apolloControlInnerSize || length > apolloControlMaximumPlain {
		return apolloControlMessage{}, errApolloControl
	}
	c.lastHost, c.hostSeen = sequence, true
	return apolloControlMessage{typeID: binary.LittleEndian.Uint16(plaintext[:2]), payload: append([]byte(nil), plaintext[4:]...)}, nil
}

func apolloControlNonce(sequence uint32, origin byte) [12]byte {
	var nonce [12]byte
	binary.LittleEndian.PutUint32(nonce[:4], sequence)
	nonce[10], nonce[11] = origin, 'C'
	return nonce
}