package crypto

import (
	"bytes"
	"encoding/base64"
	"encoding/hex"
	"fmt"

	"github.com/tjfoc/gmsm/sm3"
	"github.com/tjfoc/gmsm/sm4"
)

// SM4ECBEncrypt SM4-ECB模式加密
func SM4ECBEncrypt(plaintext []byte, key []byte) (string, error) {
	if len(key) != 16 {
		return "", fmt.Errorf("SM4密钥长度必须为16字节")
	}

	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("创建SM4 cipher失败: %v", err)
	}

	padded := pkcs7Padding(plaintext, block.BlockSize())

	ciphertext := make([]byte, len(padded))
	for i := 0; i < len(padded); i += block.BlockSize() {
		block.Encrypt(ciphertext[i:i+block.BlockSize()], padded[i:i+block.BlockSize()])
	}

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

// SM4ECBDecrypt SM4-ECB模式解密
func SM4ECBDecrypt(encryptedData string, key []byte) ([]byte, error) {
	if len(key) != 16 {
		return nil, fmt.Errorf("SM4密钥长度必须为16字节")
	}

	ciphertext, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return nil, fmt.Errorf("base64解码失败: %v", err)
	}

	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("创建SM4 cipher失败: %v", err)
	}

	blockSize := block.BlockSize()
	if len(ciphertext)%blockSize != 0 {
		return nil, fmt.Errorf("密文长度不是块大小的整数倍")
	}

	plaintext := make([]byte, len(ciphertext))
	for i := 0; i < len(ciphertext); i += blockSize {
		block.Decrypt(plaintext[i:i+blockSize], ciphertext[i:i+blockSize])
	}

	plaintext, err = pkcs7UnPadding(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}

	return plaintext, nil
}

// SM3WithSalt 带盐的SM3哈希计算
// 对应Java的 SmUtil.sm3WithSalt(saltBytes).digestHex(data)
// 实现方式：SM3(salt + data)，输出十六进制字符串
func SM3WithSalt(data string, salt string) string {
	h := sm3.New()
	h.Write([]byte(salt))
	h.Write([]byte(data))
	return hex.EncodeToString(h.Sum(nil))
}

// SM3Hash 计算SM3哈希值
func SM3Hash(data []byte) string {
	return hex.EncodeToString(sm3.Sm3Sum(data))
}

// pkcs7Padding PKCS7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// pkcs7UnPadding 去除PKCS7填充
func pkcs7UnPadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, fmt.Errorf("数据为空")
	}
	padding := int(data[length-1])
	if padding > length || padding == 0 {
		return nil, fmt.Errorf("无效的填充")
	}
	for i := length - padding; i < length; i++ {
		if data[i] != byte(padding) {
			return nil, fmt.Errorf("无效的填充")
		}
	}
	return data[:length-padding], nil
}