package sm2

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/tjfoc/gmsm/sm2"
	"math/big"
	"strings"
)

// 生成公钥、私钥
func GenerateSM2Key() (PublicKey string, PrivateKey string, err error) {
	// 生成私钥、公钥
	privKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {

		return "", "", err
	}
	return PublicKeyToString(&privKey.PublicKey), PrivateKeyToString(privKey), nil
}

// GenerateKey 生成密钥对
func GenerateKey() (string, string) {
	pri, _ := sm2.GenerateKey(rand.Reader)
	hexPri := pri.D.Text(16)
	// 获取公钥
	publicKeyHex := PublicKeyToString(&pri.PublicKey)
	return strings.ToUpper(hexPri), publicKeyHex
}

// PublicKeyToString 公钥sm2.PublicKey转字符串(与java中org.bouncycastle.crypto生成的公私钥完全互通使用)
func PublicKeyToString(publicKey *sm2.PublicKey) string {
	xBytes := publicKey.X.Bytes()
	yBytes := publicKey.Y.Bytes()

	// 确保坐标字节切片长度相同
	byteLen := len(xBytes)
	if len(yBytes) > byteLen {
		byteLen = len(yBytes)
	}

	// 为坐标补齐前导零
	xBytes = append(make([]byte, byteLen-len(xBytes)), xBytes...)
	yBytes = append(make([]byte, byteLen-len(yBytes)), yBytes...)

	// 添加 "04" 前缀
	publicKeyBytes := append([]byte{0x04}, append(xBytes, yBytes...)...)

	return strings.ToUpper(hex.EncodeToString(publicKeyBytes))
}

// PrivateKeyToString 私钥sm2.PrivateKey 转字符串(与java中org.bouncycastle.crypto生成的公私钥完全互通使用)
func PrivateKeyToString(privateKey *sm2.PrivateKey) string {
	return strings.ToUpper(hex.EncodeToString(privateKey.D.Bytes()))
}

func SM2Decrypt(cipherText, publicKey string, privateKey string) (string, error) {
	if cipherText == "" {
		return "", nil
	}
	decodedBytes, err := base64.StdEncoding.DecodeString(cipherText)
	if err != nil {

		return "", nil
	}
	decrypt, err := decryptLoc(publicKey, privateKey, string(decodedBytes))
	if err != nil {
		return "", err
	}
	return decrypt, nil
}

func SM2Encrypt(cipherText string, privateKey string) (string, error) {
	if cipherText == "" {
		return "", nil
	}
	decrypt, err := encryptLoc(privateKey, cipherText)
	if err != nil {
		return "", err
	}
	return decrypt, nil
}

func encryptLoc(publicKeyStr, data string) (string, error) {
	publicKeyObj, err := StringToPublicKey(publicKeyStr)
	if err != nil {
		return "", err
	}
	decrypt, err := sm2.Encrypt(publicKeyObj, []byte(data), rand.Reader, sm2.C1C2C3)
	if err != nil {
		return "", err
	}
	resultStr := hex.EncodeToString(decrypt)
	return base64.StdEncoding.EncodeToString([]byte(resultStr)), nil
}

func decryptLoc(publicKeyStr, privateKeyStr, cipherText string) (string, error) {
	publicKeyObj, err := StringToPublicKey(publicKeyStr)
	if err != nil {
		fmt.Println(err)
	}
	privateKeyObj, err := StringToPrivateKey(privateKeyStr, publicKeyObj)
	if err != nil {
		fmt.Println(err)
	}
	decodeString, err := hex.DecodeString(cipherText)
	decrypt, err := sm2.Decrypt(privateKeyObj, decodeString, sm2.C1C2C3)
	if err != nil {
		fmt.Println(err)
	}
	resultStr := string(decrypt)
	return resultStr, nil
}

// StringToPrivateKey 私钥还原为 sm2.PrivateKey对象(与java中org.bouncycastle.crypto生成的公私钥完全互通使用)
func StringToPrivateKey(privateKeyStr string, publicKey *sm2.PublicKey) (*sm2.PrivateKey, error) {
	privateKeyBytes, err := hex.DecodeString(privateKeyStr)
	if err != nil {
		return nil, err
	}

	// 将字节切片转换为大整数
	d := new(big.Int).SetBytes(privateKeyBytes)

	// 创建 sm2.PrivateKey 对象
	privateKey := &sm2.PrivateKey{
		PublicKey: *publicKey,
		D:         d,
	}

	return privateKey, nil
}

// StringToPublicKey 公钥字符串还原为 sm2.PublicKey 对象(与java中org.bouncycastle.crypto生成的公私钥完全互通使用)
func StringToPublicKey(publicKeyStr string) (*sm2.PublicKey, error) {
	publicKeyBytes, err := hex.DecodeString(publicKeyStr)
	if err != nil {
		return nil, err
	}

	// 提取 x 和 y 坐标字节切片
	curve := sm2.P256Sm2().Params()
	byteLen := (curve.BitSize + 7) / 8
	xBytes := publicKeyBytes[1 : byteLen+1]
	yBytes := publicKeyBytes[byteLen+1 : 2*byteLen+1]

	// 将字节切片转换为大整数
	x := new(big.Int).SetBytes(xBytes)
	y := new(big.Int).SetBytes(yBytes)

	// 创建 sm2.PublicKey 对象
	publicKey := &sm2.PublicKey{
		Curve: curve,
		X:     x,
		Y:     y,
	}

	return publicKey, nil
}

// 验证签名
func VerSm2Sig(pub *sm2.PublicKey, msg []byte, sign []byte) bool {
	isok := pub.Verify(msg, sign)
	return isok
}
