package speedport

import (
	"crypto/aes"
	"crypto/subtle"
	"encoding/binary"
	"errors"
	"fmt"
)

// CCM implements AES-CCM (Counter with CBC-MAC) as defined in RFC 3610.
// This is needed because Go's standard library does not include CCM.
type CCM struct {
	key     []byte
	tagSize int
}

// NewCCM creates a new AES-CCM instance.
// tagSize is the authentication tag length in bytes (4, 6, 8, 10, 12, 14, or 16).
func NewCCM(key []byte, tagSize int) (*CCM, error) {
	if tagSize < 4 || tagSize > 16 || tagSize%2 != 0 {
		return nil, fmt.Errorf("invalid tag size %d: must be even and between 4-16", tagSize)
	}
	switch len(key) {
	case 16, 24, 32:
		// valid AES key sizes
	default:
		return nil, fmt.Errorf("invalid key size %d", len(key))
	}
	return &CCM{key: key, tagSize: tagSize}, nil
}

// Decrypt decrypts and authenticates ciphertextWithTag using the given nonce and optional additional data.
// ciphertextWithTag = ciphertext || tag (tag is appended at the end).
func (c *CCM) Decrypt(nonce, ciphertextWithTag, additionalData []byte) ([]byte, error) {
	if len(nonce) < 7 || len(nonce) > 13 {
		return nil, fmt.Errorf("invalid nonce size %d: must be 7-13 bytes", len(nonce))
	}
	if len(ciphertextWithTag) < c.tagSize {
		return nil, errors.New("ciphertext too short for tag")
	}

	block, err := aes.NewCipher(c.key)
	if err != nil {
		return nil, fmt.Errorf("create cipher: %w", err)
	}

	// Split ciphertext and tag
	ct := ciphertextWithTag[:len(ciphertextWithTag)-c.tagSize]
	tag := ciphertextWithTag[len(ciphertextWithTag)-c.tagSize:]

	// L = 15 - len(nonce) (number of bytes for the length field)
	L := 15 - len(nonce)

	// --- CTR Decryption ---

	// Build A_0 (first counter block, used for tag encryption)
	a0 := make([]byte, 16)
	a0[0] = byte(L - 1) // flags for counter: just L-1 in lower 3 bits
	copy(a0[1:1+len(nonce)], nonce)
	// Counter bytes are 0 for A_0

	// Encrypt A_0 to get S_0
	s0 := make([]byte, 16)
	block.Encrypt(s0, a0)

	// Decrypt ciphertext using CTR starting from A_1
	plaintext := make([]byte, len(ct))
	ctr := make([]byte, 16)
	copy(ctr, a0)

	for i := 0; i < len(ct); i += aes.BlockSize {
		// Increment counter (last L bytes, big-endian)
		c.incrementCounter(ctr, L)

		var keystreamBlock [aes.BlockSize]byte
		block.Encrypt(keystreamBlock[:], ctr)

		end := i + aes.BlockSize
		if end > len(ct) {
			end = len(ct)
		}
		for j := i; j < end; j++ {
			plaintext[j] = ct[j] ^ keystreamBlock[j-i]
		}
	}

	// --- CBC-MAC Computation ---

	// Build B_0
	b0 := make([]byte, 16)
	hasAAD := len(additionalData) > 0
	aadFlag := byte(0)
	if hasAAD {
		aadFlag = 1 << 6
	}
	tagEncoding := byte((c.tagSize-2)/2) << 3
	b0[0] = aadFlag | tagEncoding | byte(L-1)
	copy(b0[1:1+len(nonce)], nonce)

	// Encode message length (plaintext length) in last L bytes, big-endian
	msgLen := len(ct)
	for i := 0; i < L; i++ {
		b0[15-i] = byte(msgLen >> (8 * uint(i)))
	}

	// Start CBC-MAC: X_1 = E(B_0)
	mac := make([]byte, 16)
	block.Encrypt(mac, b0)

	// Process AAD if present
	if hasAAD {
		aadBlock := c.formatAAD(additionalData)
		for i := 0; i < len(aadBlock); i += aes.BlockSize {
			for j := 0; j < aes.BlockSize; j++ {
				mac[j] ^= aadBlock[i+j]
			}
			block.Encrypt(mac, mac)
		}
	}

	// Process plaintext blocks
	for i := 0; i < len(plaintext); i += aes.BlockSize {
		end := i + aes.BlockSize
		if end > len(plaintext) {
			// Last partial block: XOR available bytes, rest stays as-is
			for j := i; j < len(plaintext); j++ {
				mac[j-i] ^= plaintext[j]
			}
		} else {
			for j := 0; j < aes.BlockSize; j++ {
				mac[j] ^= plaintext[i+j]
			}
		}
		block.Encrypt(mac, mac)
	}

	// Compute expected tag: T = first tagSize bytes of (MAC XOR S_0)
	expectedTag := make([]byte, c.tagSize)
	for i := 0; i < c.tagSize; i++ {
		expectedTag[i] = mac[i] ^ s0[i]
	}

	// Constant-time comparison
	if subtle.ConstantTimeCompare(expectedTag, tag) != 1 {
		return nil, errors.New("CCM authentication failed: tag mismatch")
	}

	return plaintext, nil
}

// Encrypt encrypts plaintext using the given nonce and optional additional data.
// Returns ciphertext || tag.
func (c *CCM) Encrypt(nonce, plaintext, additionalData []byte) ([]byte, error) {
	if len(nonce) < 7 || len(nonce) > 13 {
		return nil, fmt.Errorf("invalid nonce size %d: must be 7-13 bytes", len(nonce))
	}

	block, err := aes.NewCipher(c.key)
	if err != nil {
		return nil, fmt.Errorf("create cipher: %w", err)
	}

	L := 15 - len(nonce)

	// --- CBC-MAC Computation ---

	b0 := make([]byte, 16)
	hasAAD := len(additionalData) > 0
	aadFlag := byte(0)
	if hasAAD {
		aadFlag = 1 << 6
	}
	tagEncoding := byte((c.tagSize-2)/2) << 3
	b0[0] = aadFlag | tagEncoding | byte(L-1)
	copy(b0[1:1+len(nonce)], nonce)

	msgLen := len(plaintext)
	for i := 0; i < L; i++ {
		b0[15-i] = byte(msgLen >> (8 * uint(i)))
	}

	mac := make([]byte, 16)
	block.Encrypt(mac, b0)

	if hasAAD {
		aadBlock := c.formatAAD(additionalData)
		for i := 0; i < len(aadBlock); i += aes.BlockSize {
			for j := 0; j < aes.BlockSize; j++ {
				mac[j] ^= aadBlock[i+j]
			}
			block.Encrypt(mac, mac)
		}
	}

	for i := 0; i < len(plaintext); i += aes.BlockSize {
		end := i + aes.BlockSize
		if end > len(plaintext) {
			for j := i; j < len(plaintext); j++ {
				mac[j-i] ^= plaintext[j]
			}
		} else {
			for j := 0; j < aes.BlockSize; j++ {
				mac[j] ^= plaintext[i+j]
			}
		}
		block.Encrypt(mac, mac)
	}

	// --- CTR Encryption ---

	a0 := make([]byte, 16)
	a0[0] = byte(L - 1)
	copy(a0[1:1+len(nonce)], nonce)

	s0 := make([]byte, 16)
	block.Encrypt(s0, a0)

	ciphertext := make([]byte, len(plaintext))
	ctr := make([]byte, 16)
	copy(ctr, a0)

	for i := 0; i < len(plaintext); i += aes.BlockSize {
		c.incrementCounter(ctr, L)

		var keystreamBlock [aes.BlockSize]byte
		block.Encrypt(keystreamBlock[:], ctr)

		end := i + aes.BlockSize
		if end > len(plaintext) {
			end = len(plaintext)
		}
		for j := i; j < end; j++ {
			ciphertext[j] = plaintext[j] ^ keystreamBlock[j-i]
		}
	}

	// Compute tag
	tag := make([]byte, c.tagSize)
	for i := 0; i < c.tagSize; i++ {
		tag[i] = mac[i] ^ s0[i]
	}

	return append(ciphertext, tag...), nil
}

// formatAAD formats additional authenticated data per RFC 3610 Section 2.2.
func (c *CCM) formatAAD(aad []byte) []byte {
	var encoded []byte
	aadLen := len(aad)

	if aadLen < 0xFF00 {
		encoded = make([]byte, 2)
		binary.BigEndian.PutUint16(encoded, uint16(aadLen))
	} else if aadLen < (1 << 32) {
		encoded = make([]byte, 6)
		encoded[0] = 0xFF
		encoded[1] = 0xFE
		binary.BigEndian.PutUint32(encoded[2:], uint32(aadLen))
	}

	encoded = append(encoded, aad...)

	// Pad to block boundary
	if rem := len(encoded) % aes.BlockSize; rem != 0 {
		encoded = append(encoded, make([]byte, aes.BlockSize-rem)...)
	}

	return encoded
}

// incrementCounter increments the last L bytes of the counter block (big-endian).
func (c *CCM) incrementCounter(ctr []byte, L int) {
	for i := 15; i >= 16-L; i-- {
		ctr[i]++
		if ctr[i] != 0 {
			return
		}
	}
}