package cmb

import (
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"github.com/tjfoc/gmsm/sm4"
)

func generateSm4Key() (string, error) {
	keyBytes, err := generateKey()
	if err != nil {
		return "", err
	}
	return ByteArrayToHexString(keyBytes)
}

func generateKey() ([]byte, error) {
	key := make([]byte, sm4.BlockSize)
	_, err := rand.Read(key)
	if err != nil {
		return nil, err
	}
	return key, nil
}

// ByteArrayToHexString 将字节数组转换为十六进制字符串
func ByteArrayToHexString(byteArray []byte) (string, error) {
	if byteArray == nil {
		return "", fmt.Errorf("byteArray is nil")
	}
	result := make([]byte, 0, len(byteArray)*2)
	for _, b := range byteArray {
		// 处理高 4 位
		result = append(result, byteToHexChar(b>>4))
		// 处理低 4 位
		result = append(result, byteToHexChar(b&0x0F))
	}
	return string(result), nil
}

// byteToHexChar 将一个 4 位字节转换为对应的十六进制字符
func byteToHexChar(n byte) byte {
	if n < 10 {
		return '0' + n
	}
	return 'a' + n - 10
}

func SM4Encrypt(contentBytes []byte, encryptKey string) (string, error) {
	d, err := base64.StdEncoding.DecodeString(encryptKey)
	if err != nil {
		return "", err
	}

	cipherBlock, err := sm4.NewCipher(d)
	if err != nil {
		return "", err
	}

	blockSize := cipherBlock.BlockSize()
	iv := make([]byte, blockSize)
	for i := 0; i < blockSize; i++ {
		iv[i] = 0
	}

	blockMode := cipher.NewCBCEncrypter(cipherBlock, iv)

	padding := blockSize - len(contentBytes)%blockSize
	for i := 0; i < padding; i++ {
		contentBytes = append(contentBytes, byte(padding))
	}
	cipherText := make([]byte, len(contentBytes))
	blockMode.CryptBlocks(cipherText, contentBytes)

	return base64.StdEncoding.EncodeToString(cipherText), nil
}

func SM4Decrypt(encrypted, encryptKey string) (string, error) {
	d, err := base64.StdEncoding.DecodeString(encryptKey)
	if err != nil {
		return "", err
	}

	cipherBlock, err := sm4.NewCipher(d)
	if err != nil {
		return "", err
	}

	blockSize := cipherBlock.BlockSize()
	iv := make([]byte, blockSize)
	for i := 0; i < blockSize; i++ {
		iv[i] = 0
	}

	cipherText, err := base64.StdEncoding.DecodeString(encrypted)
	if err != nil {
		return "", err
	}

	plainText := make([]byte, len(cipherText))
	blockMode := cipher.NewCBCDecrypter(cipherBlock, iv)
	blockMode.CryptBlocks(plainText, cipherText)

	plainTextLen := len(plainText)
	padding := int(plainText[plainTextLen-1])
	buff := plainText[:plainTextLen-padding]

	return string(buff), nil
}