package crypto

import (
	"bytes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"

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

var (
	sm3Salt []byte
	sm4Key  []byte
)

func InitCrypto(salt, key string) {
	sm3Salt = []byte(salt)
	keyBytes := []byte(key)
	if len(keyBytes) < 16 {
		padded := make([]byte, 16)
		copy(padded, keyBytes)
		sm4Key = padded
	} else if len(keyBytes) > 16 {
		sm4Key = keyBytes[:16]
	} else {
		sm4Key = keyBytes
	}
}

func SM3Hash(data []byte) string {
	h := sm3.New()
	h.Write(data)
	hashBytes := h.Sum(nil)
	return hex.EncodeToString(hashBytes)
}

func SM3HashWithSalt(data string) string {
	h := hmac.New(sm3.New, sm3Salt)
	h.Write([]byte(data))
	return hex.EncodeToString(h.Sum(nil))
}

func GenerateSign(timestamp, encryptedData string) string {
	signStr := timestamp + encryptedData
	return SM3HashWithSalt(signStr)
}

func VerifySign(timestamp, encryptedData, sign string) bool {
	expectedSign := GenerateSign(timestamp, encryptedData)
	return hmac.Equal([]byte(expectedSign), []byte(sign))
}

func SM4CBCEncrypt(plaintext []byte) (string, error) {
	if len(sm4Key) != 16 {
		return "", fmt.Errorf("SM4密钥长度必须为16字节")
	}

	block, err := sm4.NewCipher(sm4Key)
	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
}

func SM4CBCDecrypt(encryptedData string) ([]byte, error) {
	if len(sm4Key) != 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(sm4Key)
	if err != nil {
		return nil, fmt.Errorf("创建SM4 cipher失败: %v", err)
	}

	blockSize := block.BlockSize()
	if len(data) < blockSize {
		return nil, errors.New("数据长度不足")
	}
	iv := data[:blockSize]
	ciphertext := data[blockSize:]

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

	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
}

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

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