package algos

// SM3 SM3哈希算法实现
type SM3 struct {
	IV []uint32
	TJ []uint32
}

// NewSM3 创建SM3实例
func NewSM3() *SM3 {
	return &SM3{
		IV: []uint32{1937774191, 1226093241, 388252375, 3666478592, 2842636476, 372324522, 3817729613, 2969243214},
		TJ: []uint32{
			2043430169, 2043430169, 2043430169, 2043430169, 2043430169, 2043430169, 2043430169, 2043430169,
			2043430169, 2043430169, 2043430169, 2043430169, 2043430169, 2043430169, 2043430169, 2043430169,
			2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042,
			2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042,
			2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042,
			2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042,
			2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042,
			2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042, 2055708042,
		},
	}
}

// rotateLeft 循环左移
func (sm3 *SM3) rotateLeft(a uint32, k int) uint32 {
	k = k % 32
	return ((a << uint(k)) & 0xFFFFFFFF) | ((a & 0xFFFFFFFF) >> uint(32-k))
}

// FFJ FF函数
func (sm3 *SM3) FFJ(X, Y, Z uint32, j int) uint32 {
	if j >= 0 && j < 16 {
		return X ^ Y ^ Z
	} else if j >= 16 && j < 64 {
		return (X & Y) | (X & Z) | (Y & Z)
	}
	return 0
}

// GGJ GG函数
func (sm3 *SM3) GGJ(X, Y, Z uint32, j int) uint32 {
	if j >= 0 && j < 16 {
		return X ^ Y ^ Z
	} else if j >= 16 && j < 64 {
		return (X & Y) | ((^X) & Z)
	}
	return 0
}

// P0 压缩函数P0
func (sm3 *SM3) P0(X uint32) uint32 {
	return X ^ sm3.rotateLeft(X, 9) ^ sm3.rotateLeft(X, 17)
}

// P1 压缩函数P1
func (sm3 *SM3) P1(X uint32) uint32 {
	return X ^ sm3.rotateLeft(X, 15) ^ sm3.rotateLeft(X, 23)
}

// CF 压缩函数
func (sm3 *SM3) CF(Vi []uint32, Bi []byte) []uint32 {
	W := make([]uint32, 68)
	for i := 0; i < 16; i++ {
		weight := uint32(0x1000000)
		data := uint32(0)
		for k := i * 4; k < (i+1)*4; k++ {
			data = data + uint32(Bi[k])*weight
			weight = weight / 0x100
		}
		W[i] = data
	}

	for j := 16; j < 68; j++ {
		W[j] = sm3.P1(W[j-16]^W[j-9]^sm3.rotateLeft(W[j-3], 15)) ^ sm3.rotateLeft(W[j-13], 7) ^ W[j-6]
	}

	W1 := make([]uint32, 64)
	for j := 0; j < 64; j++ {
		W1[j] = W[j] ^ W[j+4]
	}

	A, B, C, D, E, F, G, H := Vi[0], Vi[1], Vi[2], Vi[3], Vi[4], Vi[5], Vi[6], Vi[7]

	for j := 0; j < 64; j++ {
		SS1 := sm3.rotateLeft((sm3.rotateLeft(A, 12)+E+sm3.rotateLeft(sm3.TJ[j], j%32))&0xFFFFFFFF, 7)
		SS2 := SS1 ^ sm3.rotateLeft(A, 12)
		TT1 := (sm3.FFJ(A, B, C, j) + D + SS2 + W1[j]) & 0xFFFFFFFF
		TT2 := (sm3.GGJ(E, F, G, j) + H + SS1 + W[j]) & 0xFFFFFFFF
		D = C
		C = sm3.rotateLeft(B, 9)
		B = A
		A = TT1
		H = G
		G = sm3.rotateLeft(F, 19)
		F = E
		E = sm3.P0(TT2)
	}

	V := make([]uint32, 8)
	V[0] = A ^ Vi[0]
	V[1] = B ^ Vi[1]
	V[2] = C ^ Vi[2]
	V[3] = D ^ Vi[3]
	V[4] = E ^ Vi[4]
	V[5] = F ^ Vi[5]
	V[6] = G ^ Vi[6]
	V[7] = H ^ Vi[7]

	return V
}

// Hash 计算SM3哈希
func (sm3 *SM3) Hash(msg []byte) []byte {
	msgLen := len(msg)
	reserve1 := msgLen % 64
	paddingLen := 0
	if reserve1 > 55 {
		paddingLen = 64 - reserve1 + 64
	} else {
		paddingLen = 64 - reserve1
	}

	msgPadding := make([]byte, msgLen+paddingLen)
	copy(msgPadding, msg)
	msgPadding[msgLen] = 0x80

	lengthBits := uint64(msgLen) * 8
	for i := 0; i < 8; i++ {
		msgPadding[len(msgPadding)-1-i] = byte(lengthBits >> uint(8*i))
	}

	groupCount := len(msgPadding) / 64

	V := make([]uint32, len(sm3.IV))
	copy(V, sm3.IV)

	for i := 0; i < groupCount; i++ {
		B := msgPadding[i*64 : (i+1)*64]
		V = sm3.CF(V, B)
	}

	result := make([]byte, 32)
	for i := 0; i < 8; i++ {
		result[i*4] = byte((V[i] >> 24) & 0xFF)
		result[i*4+1] = byte((V[i] >> 16) & 0xFF)
		result[i*4+2] = byte((V[i] >> 8) & 0xFF)
		result[i*4+3] = byte(V[i] & 0xFF)
	}

	return result
}