package crypto

import (
	"bytes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"math/big"
	"time"

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

// ============================================================
// SM4-CBC 加密/解密
// ============================================================

// SM4CBCEncrypt SM4-CBC模式加密
func SM4CBCEncrypt(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())
	iv := make([]byte, block.BlockSize())
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", fmt.Errorf("生成IV失败: %v", err)
	}
	mode := cipher.NewCBCEncrypter(block, iv)
	ciphertext := make([]byte, len(padded))
	mode.CryptBlocks(ciphertext, padded)
	result := append(iv, ciphertext...)
	return base64.StdEncoding.EncodeToString(result), nil
}

// SM4CBCDecrypt SM4-CBC模式解密
func SM4CBCDecrypt(encryptedData string, key []byte) ([]byte, error) {
	if len(key) != 16 {
		return nil, fmt.Errorf("SM4密钥长度必须为16字节")
	}
	data, 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(data) < blockSize {
		return nil, fmt.Errorf("数据长度不足")
	}
	iv := data[:blockSize]
	ciphertext := data[blockSize:]
	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)
	plaintext, err = pkcs7UnPadding(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}
	return plaintext, nil
}

// 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
}

// ============================================================
// SM3 哈希 / HMAC-SM3
// ============================================================

// SM3Hash 计算SM3哈希值
func SM3Hash(data []byte, encoding string) (string, error) {
	h := sm3.New()
	h.Write(data)
	hashBytes := h.Sum(nil)
	if encoding == "base64" {
		return base64.StdEncoding.EncodeToString(hashBytes), nil
	}
	return hex.EncodeToString(hashBytes), nil
}

// SM3HashString 计算字符串的SM3哈希值
func SM3HashString(data string, encoding string) (string, error) {
	return SM3Hash([]byte(data), encoding)
}

// HMACSM3 HMAC-SM3计算
// 对应 Java hutool 的 sm3WithSalt(saltBytes).digestHex(data)
func HMACSM3(data []byte, key []byte) string {
	h := hmac.New(sm3.New, key)
	h.Write(data)
	return hex.EncodeToString(h.Sum(nil))
}

// ============================================================
// 时间戳 / 随机数
// ============================================================

// GenerateTimestampMillis 生成毫秒级时间戳字符串
func GenerateTimestampMillis() string {
	return fmt.Sprintf("%d", time.Now().UnixMilli())
}

// GenerateNonce 生成指定长度的随机字符串
func GenerateNonce(length int) (string, error) {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, length)
	for i := range b {
		num, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
		if err != nil {
			return "", err
		}
		b[i] = charset[num.Int64()]
	}
	return string(b), nil
}

// ============================================================
// 签名生成与验证（文档专用）
// ============================================================

// GenerateSign 生成签名
// 规则：HMAC-SM3(salt, timestamp + encryptedData) -> hex
func GenerateSign(timestamp string, encryptedData string, sm3Salt []byte) string {
	data := timestamp + encryptedData
	return HMACSM3([]byte(data), sm3Salt)
}

// VerifySign 验证签名
func VerifySign(timestamp string, encryptedData string, sign string, sm3Salt []byte) bool {
	expected := GenerateSign(timestamp, encryptedData, sm3Salt)
	return expected == sign
}