package protocol

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

// OMAC1 (CMAC) Implementation for AES-128
// RFC 4493

func shiftLeft(in []byte, out []byte) {
	var carry byte
	for i := 15; i >= 0; i-- {
		b := in[i]
		out[i] = (b << 1) | carry
		carry = (b >> 7) & 1
	}
}

func xorBlock(a, b []byte) {
	for i := 0; i < 16; i++ {
		a[i] ^= b[i]
	}
}

func generateSubkeys(cipherBlock cipher.Block) (k1, k2 []byte) {
	l := make([]byte, 16)
	cipherBlock.Encrypt(l, make([]byte, 16))

	k1 = make([]byte, 16)
	shiftLeft(l, k1)
	if l[0]&0x80 != 0 {
		k1[15] ^= 0x87
	}

	k2 = make([]byte, 16)
	shiftLeft(k1, k2)
	if k1[0]&0x80 != 0 {
		k2[15] ^= 0x87
	}
	return
}

func omac1(cipherBlock cipher.Block, data []byte) []byte {
	k1, k2 := generateSubkeys(cipherBlock)

	n := len(data)
	numBlocks := (n + 15) / 16
	if numBlocks == 0 {
		numBlocks = 1
	}

	lastBlock := make([]byte, 16)
	flagComplete := (n > 0 && n%16 == 0)

	if flagComplete {
		copy(lastBlock, data[n-16:])
		xorBlock(lastBlock, k1)
	} else {
		remaining := n % 16
		copy(lastBlock, data[n-remaining:])
		lastBlock[remaining] = 0x80 // padding
		xorBlock(lastBlock, k2)
	}

	x := make([]byte, 16)
	y := make([]byte, 16)

	for i := 0; i < numBlocks-1; i++ {
		copy(y, data[i*16:(i+1)*16])
		xorBlock(y, x)
		cipherBlock.Encrypt(x, y)
	}

	xorBlock(x, lastBlock)
	cipherBlock.Encrypt(x, x)
	return x
}

// EAX implementation functions

// omacWithTag computes OMAC(t || m)
func omacWithTag(block cipher.Block, tag byte, m []byte) []byte {
	// EAX uses [t] as prefix for OMAC
	// but standard OMAC implies it's just the message.
	// EAX spec: OMAC^t_K(M) = OMAC_K ([t]_n || M)
	// where n is block size (16).

	// We construct a new slice with the prefix
	buf := make([]byte, 16+len(m))
	buf[15] = tag // The last byte of the first block is the tag?
	// Wait, EAX spec: [t]_n is the byte t padded with zeros to block size n.
	// Usually it implies the last byte is t, or first?
	// "Let [t]_n denote the n-byte string that consists of n-1 zero bytes followed by the byte t."
	// So for n=16, 15 zeros then t.

	copy(buf[16:], m)
	return omac1(block, buf)
}

// EncryptEAX performs EAX encryption and returns ciphertext + mac
// Key and Nonce must be valid for AES-128 (16 bytes).
func EncryptEAX(key, nonce, header, data []byte) (ciphertext []byte, mac []byte, err error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, nil, err
	}

	// 1. Calculate Nonce MAC (N)
	// N = OMAC^0_K(Nonce)
	nMac := omacWithTag(block, 0, nonce)

	// 2. Calculate Header MAC (H)
	// H = OMAC^1_K(Header)
	hMac := omacWithTag(block, 1, header)

	// 3. Encrypt Data (C)
	// CTR using N as counter/IV
	ciphertext = make([]byte, len(data))
	ctrStream := cipher.NewCTR(block, nMac)
	ctrStream.XORKeyStream(ciphertext, data)

	// 4. Calculate Ciphertext MAC (C_MAC)
	// C_MAC = OMAC^2_K(Ciphertext)
	cMac := omacWithTag(block, 2, ciphertext)

	// 5. Final Tag = N ^ H ^ C_MAC
	tag := make([]byte, 16)
	for i := 0; i < 16; i++ {
		tag[i] = nMac[i] ^ hMac[i] ^ cMac[i]
	}

	// TS3 uses 8 bytes of the MAC
	return ciphertext, tag[:8], nil
}

// DecryptEAX verifies and decrypts
func DecryptEAX(key, nonce, header, data, expectedMac []byte) (plaintext []byte, err error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 1. N = OMAC^0_K(Nonce)
	nMac := omacWithTag(block, 0, nonce)

	// 2. H = OMAC^1_K(Header)
	hMac := omacWithTag(block, 1, header)

	// 3. C_MAC = OMAC^2_K(Data is Ciphertext here)
	cMac := omacWithTag(block, 2, data)

	// 4. Calculate Tag
	tag := make([]byte, 16)
	for i := 0; i < 16; i++ {
		tag[i] = nMac[i] ^ hMac[i] ^ cMac[i]
	}

	// Verify MAC
	if len(expectedMac) > 16 {
		return nil, errors.New("expected MAC too long")
	}

	// Constant time compare needed ideally
	for i := 0; i < len(expectedMac); i++ {
		if tag[i] != expectedMac[i] {
			return nil, errors.New("mac verification failed")
		}
	}

	// 5. Decrypt
	plaintext = make([]byte, len(data))
	ctrStream := cipher.NewCTR(block, nMac)
	ctrStream.XORKeyStream(plaintext, data)

	return plaintext, nil
}