package sm4

// [GM/T] SM4 GB/T 32907-2016

import (
	"encoding/binary"
	"math/bits"
)

type convert func(uint32) uint32

// Encrypt one block from src into dst, using the expanded key xk.
func encryptBlockGo(xk []uint32, dst, src []byte) {
	_ = src[15]    // early bounds check
	dst = dst[:16] // early bounds check
	_ = xk[31]     // bounds check elimination hint

	var b0, b1, b2, b3 uint32
	b0 = binary.BigEndian.Uint32(src[0:4])
	b1 = binary.BigEndian.Uint32(src[4:8])
	b2 = binary.BigEndian.Uint32(src[8:12])
	b3 = binary.BigEndian.Uint32(src[12:16])

	b0 ^= t(b1 ^ b2 ^ b3 ^ xk[0])
	b1 ^= t(b2 ^ b3 ^ b0 ^ xk[1])
	b2 ^= t(b3 ^ b0 ^ b1 ^ xk[2])
	b3 ^= t(b0 ^ b1 ^ b2 ^ xk[3])

	b0 ^= precompute_t(b1 ^ b2 ^ b3 ^ xk[4])
	b1 ^= precompute_t(b2 ^ b3 ^ b0 ^ xk[5])
	b2 ^= precompute_t(b3 ^ b0 ^ b1 ^ xk[6])
	b3 ^= precompute_t(b0 ^ b1 ^ b2 ^ xk[7])

	b0 ^= precompute_t(b1 ^ b2 ^ b3 ^ xk[8])
	b1 ^= precompute_t(b2 ^ b3 ^ b0 ^ xk[9])
	b2 ^= precompute_t(b3 ^ b0 ^ b1 ^ xk[10])
	b3 ^= precompute_t(b0 ^ b1 ^ b2 ^ xk[11])

	b0 ^= precompute_t(b1 ^ b2 ^ b3 ^ xk[12])
	b1 ^= precompute_t(b2 ^ b3 ^ b0 ^ xk[13])
	b2 ^= precompute_t(b3 ^ b0 ^ b1 ^ xk[14])
	b3 ^= precompute_t(b0 ^ b1 ^ b2 ^ xk[15])

	b0 ^= precompute_t(b1 ^ b2 ^ b3 ^ xk[16])
	b1 ^= precompute_t(b2 ^ b3 ^ b0 ^ xk[17])
	b2 ^= precompute_t(b3 ^ b0 ^ b1 ^ xk[18])
	b3 ^= precompute_t(b0 ^ b1 ^ b2 ^ xk[19])

	b0 ^= precompute_t(b1 ^ b2 ^ b3 ^ xk[20])
	b1 ^= precompute_t(b2 ^ b3 ^ b0 ^ xk[21])
	b2 ^= precompute_t(b3 ^ b0 ^ b1 ^ xk[22])
	b3 ^= precompute_t(b0 ^ b1 ^ b2 ^ xk[23])

	b0 ^= precompute_t(b1 ^ b2 ^ b3 ^ xk[24])
	b1 ^= precompute_t(b2 ^ b3 ^ b0 ^ xk[25])
	b2 ^= precompute_t(b3 ^ b0 ^ b1 ^ xk[26])
	b3 ^= precompute_t(b0 ^ b1 ^ b2 ^ xk[27])

	b0 ^= t(b1 ^ b2 ^ b3 ^ xk[28])
	b1 ^= t(b2 ^ b3 ^ b0 ^ xk[29])
	b2 ^= t(b3 ^ b0 ^ b1 ^ xk[30])
	b3 ^= t(b0 ^ b1 ^ b2 ^ xk[31])

	binary.BigEndian.PutUint32(dst[:], b3)
	binary.BigEndian.PutUint32(dst[4:], b2)
	binary.BigEndian.PutUint32(dst[8:], b1)
	binary.BigEndian.PutUint32(dst[12:], b0)
}

// Key expansion algorithm.
func expandKeyGo(key []byte, enc, dec []uint32) {
	// Encryption key setup.
	enc = enc[:rounds]
	var i int
	var mk [4]uint32
	var k [rounds + 4]uint32

	key = key[:KeySize]
	mk[0] = binary.BigEndian.Uint32(key)
	k[0] = mk[0] ^ fk[0]
	mk[1] = binary.BigEndian.Uint32(key[4:])
	k[1] = mk[1] ^ fk[1]
	mk[2] = binary.BigEndian.Uint32(key[8:])
	k[2] = mk[2] ^ fk[2]
	mk[3] = binary.BigEndian.Uint32(key[12:])
	k[3] = mk[3] ^ fk[3]

	for i = 0; i < rounds; i++ {
		k[i+4] = k[i] ^ t2(k[i+1]^k[i+2]^k[i+3]^ck[i])
		enc[i] = k[i+4]
	}

	// Derive decryption key from encryption key.
	if dec == nil {
		return
	}

	dec = dec[:rounds]
	for i = 0; i < rounds; i++ {
		dec[i] = enc[rounds-1-i]
	}
}

// Decrypt one block from src into dst, using the expanded key xk.
func decryptBlockGo(xk []uint32, dst, src []byte) {
	encryptBlockGo(xk, dst, src)
}

// L(B)
func l(b uint32) uint32 {
	return b ^ bits.RotateLeft32(b, 2) ^ bits.RotateLeft32(b, 10) ^ bits.RotateLeft32(b, 18) ^ bits.RotateLeft32(b, 24)
}

// L'(B)
func l2(b uint32) uint32 {
	return b ^ bits.RotateLeft32(b, 13) ^ bits.RotateLeft32(b, 23)
}

func _t(in uint32, fn convert) uint32 {
	var bytes [4]byte
	binary.BigEndian.PutUint32(bytes[:], in)
	for i := 0; i < 4; i++ {
		bytes[i] = sbox[bytes[i]]
	}
	return fn(binary.BigEndian.Uint32(bytes[:]))
}

// T
func t(in uint32) uint32 {
	return _t(in, l)
}

// T'
func t2(in uint32) uint32 {
	return _t(in, l2)
}

func precompute_t(in uint32) uint32 {
	return sbox_t0[byte(in>>24)] ^
		sbox_t1[byte(in>>16)] ^
		sbox_t2[byte(in>>8)] ^
		sbox_t3[byte(in)]
}