package base

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"

	"github.com/liangdas/mqant/log"
	"github.com/wenzhenxi/gorsa"
)

// 使用SHA256withRSA算法对数据进行签名
func SignWithSHA256(data []byte, privateKey string) (string, error) {
	// 解析私钥
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		return "", errors.New("invalid private key")
	}
	priKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	// 计算数据的SHA256哈希值
	hash := sha256.Sum256(data)

	// 对哈希值进行签名
	signature, err := rsa.SignPKCS1v15(rand.Reader, priKey, crypto.SHA256, hash[:])
	if err != nil {
		return "", err
	}

	// 将签名结果进行Base64编码
	return base64.StdEncoding.EncodeToString(signature), nil
}

// 使用SHA256withRSA算法对数据进行验签
func VerifyWithSHA256(data []byte, signature string, publicKey string) (bool, error) {
	// 解析公钥
	block, _ := pem.Decode([]byte(publicKey))
	if block == nil {
		return false, errors.New("invalid public key")
	}
	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return false, err
	}

	// 计算数据的SHA256哈希值
	hash := sha256.Sum256(data)

	// 对签名进行Base64解码
	signatureBytes, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return false, err
	}

	// 进行验签操作
	err = rsa.VerifyPKCS1v15(pubKey.(*rsa.PublicKey), crypto.SHA256, hash[:], signatureBytes)
	if err == nil {
		return true, nil
	} else {
		return false, err
	}
}

// 使用SHA512withRSA算法对数据进行签名
func SignWithSHA512(data []byte, privateKey string) (string, error) {
	// 解析私钥
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		return "", errors.New("invalid private key")
	}
	priKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	// 计算数据的SHA256哈希值
	hash := sha512.Sum512(data)

	// 对哈希值进行签名
	signature, err := rsa.SignPKCS1v15(rand.Reader, priKey, crypto.SHA512, hash[:])
	if err != nil {
		return "", err
	}

	// 将签名结果进行Base64编码
	return base64.StdEncoding.EncodeToString(signature), nil
}

// 使用SHA512withRSA算法对数据进行验签
func VerifyWithSHA512(data []byte, signature string, publicKey string) (bool, error) {
	// 解析公钥
	block, _ := pem.Decode([]byte(publicKey))
	if block == nil {
		return false, errors.New("invalid public key")
	}
	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return false, err
	}

	// 计算数据的SHA512哈希值
	hash := sha512.Sum512(data)

	// 对签名进行Base64解码
	signatureBytes, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return false, err
	}

	// 进行验签操作
	err = rsa.VerifyPKCS1v15(pubKey.(*rsa.PublicKey), crypto.SHA512, hash[:], signatureBytes)
	if err == nil {
		return true, nil
	} else {
		return false, err
	}
}

// 私钥签名
func RsaSignSha1(data []byte, privateKey string) ([]byte, error) {
	h := sha1.New()
	h.Write(data)
	hashed := h.Sum(nil)
	// 获取私钥
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		return nil, errors.New("private key error")
	}
	// 解析PKCS1格式的私钥
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA1, hashed)
}

// 公钥验证
func RsaSignVer(data []byte, signature, publicKey string) error {
	sign, _ := base64.StdEncoding.DecodeString(signature)
	hashed := sha1.Sum(data)
	block, _ := pem.Decode([]byte(publicKey))
	// 解析公钥
	pubInterface, _ := x509.ParsePKIXPublicKey(block.Bytes)

	// 类型断言
	pub := pubInterface.(*rsa.PublicKey)
	// 验证签名
	return rsa.VerifyPKCS1v15(pub, crypto.SHA1, hashed[:], sign)
}

func RsaPublicDecrypt(pubKey *rsa.PublicKey, data []byte) []byte {
	c := new(big.Int)
	m := new(big.Int)
	m.SetBytes(data)
	e := big.NewInt(int64(pubKey.E))
	c.Exp(m, e, pubKey.N)
	out := c.Bytes()
	skip := 0
	for i := 2; i < len(out); i++ {
		if i+1 >= len(out) {
			break
		}
		if out[i] == 0xff && out[i+1] == 0 {
			skip = i + 2
			break
		}
	}
	return out[skip:]
}

// 公钥解密
func RsaDecode(signature, publicKey string) error {
	block, _ := pem.Decode([]byte(publicKey))
	public, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		log.Error("加载公钥错误 %+v\n", err)
		return err
	}
	decodeBytes, _ := base64.StdEncoding.DecodeString(signature)
	var result string
	// 因为加密解密用rsa有长度限值，因此这个地方就看具体限值的多少了
	for len(decodeBytes) > 0 {
		decodePart := decodeBytes[:128]
		plain := RsaPublicDecrypt(public.(*rsa.PublicKey), decodePart)
		result += string(plain)
		decodeBytes = decodeBytes[128:]
	}
	log.Debug("%+v", result)
	if len(result) > 0 {
		return nil
	}
	return errors.New("check sign fail")
}

func SignRSA(prvkey []byte, hash crypto.Hash, data []byte) (string, error) {
	block, _ := pem.Decode(prvkey)
	if block == nil {
		return "", errors.New("decode private key error")
	}
	privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	// MD5 and SHA1 are not supported as they are not secure.
	var hashed []byte
	switch hash {
	case crypto.SHA224:
		h := sha256.Sum224(data)
		hashed = h[:]
	case crypto.SHA256:
		h := sha256.Sum256(data)
		hashed = h[:]
	case crypto.SHA384:
		h := sha512.Sum384(data)
		hashed = h[:]
	case crypto.SHA512:
		h := sha512.Sum512(data)
		hashed = h[:]
	}
	signature, _ := rsa.SignPKCS1v15(rand.Reader, privateKey.(*rsa.PrivateKey), hash, hashed)
	return base64.StdEncoding.EncodeToString(signature), nil
}

func RsaVerify(pubkey []byte, hash crypto.Hash, data, sig []byte) error {
	block, _ := pem.Decode(pubkey)
	if block == nil {
		return errors.New("decode public key error")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return err
	}
	// SHA1 and MD5 are not supported as they are not secure.
	var hashed []byte
	switch hash {
	case crypto.SHA224:
		h := sha256.Sum224(data)
		hashed = h[:]
	case crypto.SHA256:
		h := sha256.Sum256(data)
		hashed = h[:]
	case crypto.SHA384:
		h := sha512.Sum384(data)
		hashed = h[:]
	case crypto.SHA512:
		h := sha512.Sum512(data)
		hashed = h[:]
	}
	return rsa.VerifyPKCS1v15(pub.(*rsa.PublicKey), hash, hashed, sig)
}

func SignRSADecode(data string, publicKey string) (string, error) {
	dataBs, err := base64.RawURLEncoding.DecodeString(data)
	if err != nil {
		return "", err
	}

	gRsa := gorsa.RSASecurity{}
	if err := gRsa.SetPublicKey(publicKey); err != nil {
		return "", err
	}

	rsaData, err := gRsa.PubKeyDECRYPT(dataBs)
	if err != nil {
		return "", err
	}

	return string(rsaData), nil
}

// 解析 PKCS#8 格式的私钥
func parsePrivateKey(pemStr string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil || block.Type != "PRIVATE KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing private key")
	}
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	privateKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA private key")
	}
	return privateKey, nil
}

// 解析公钥
func parsePublicKey(pemStr string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemStr))
	if block == nil || block.Type != "PUBLIC KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing public key")
	}
	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	publicKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA public key")
	}
	return publicKey, nil
}