package tools

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"
)

// RSAGenerateKeys 生成RSA公私钥对，并返回私钥和公钥
func RSAGenerateKeys(bits int) (*rsa.PrivateKey, *rsa.PublicKey, error) {
	// 生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, nil, err
	}
	// 从私钥中提取公钥
	publicKey := &privateKey.PublicKey
	return privateKey, publicKey, nil
}

// RSASavePrivatePEMKey 将私钥保存到文件中
func RSASavePrivatePEMKey(filename string, key *rsa.PrivateKey) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	var privateKeyPEM = &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}

	return pem.Encode(outFile, privateKeyPEM)
}

// RSASavePublicPEMKey 将公钥保存到文件中
func RSASavePublicPEMKey(filename string, pubkey *rsa.PublicKey) error {
	pubASN1, err := x509.MarshalPKIXPublicKey(pubkey)
	if err != nil {
		return err
	}

	var pubKeyPEM = &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubASN1,
	}

	outFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer outFile.Close()

	return pem.Encode(outFile, pubKeyPEM)
}

// RSALoadPrivatePEMKey 从PEM文件中加载私钥
func RSALoadPrivatePEMKey(fileName string) (*rsa.PrivateKey, error) {
	privBytes, err := os.ReadFile(fileName)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(privBytes)
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing private key")
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	return privateKey, nil
}

// RSALoadPublicPEMKey 从PEM文件中加载公钥
func RSALoadPublicPEMKey(fileName string) (*rsa.PublicKey, error) {
	pubBytes, err := os.ReadFile(fileName)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(pubBytes)
	if block == nil || block.Type != "RSA PUBLIC KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing public key")
	}

	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

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

	return publicKey, nil
}

// RSAEncrypt 使用公钥加密数据
func RSAEncrypt(publicKey *rsa.PublicKey, data []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, publicKey, data)
}

// RSADecrypt 使用私钥解密数据
func RSADecrypt(privateKey *rsa.PrivateKey, ciphertext []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, privateKey, ciphertext)
}