package rsa

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
)

var _ Public = (*rsaPub)(nil)
var _ Private = (*rsaPri)(nil)

type Public interface {
	i()
	EncryptURLEncoding(encryptStr string) (string, error)
	Encrypt(encryptStr string) (string, error)
}

type Private interface {
	i()
	Decrypt(decryptStr string) (string, error)
	DecryptURLEncoding(decryptStr string) (string, error)
}

type rsaPub struct {
	PublicKey string
}

type rsaPri struct {
	PrivateKey string
}

func NewPublic(publicKey string) Public {
	return &rsaPub{
		PublicKey: publicKey,
	}
}

func NewPrivate(privateKey string) Private {
	return &rsaPri{
		PrivateKey: privateKey,
	}
}

func (pub *rsaPub) i() {}

// parsePublicKey 解析多种格式的公钥
func parsePublicKey(pemData []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pemData)
	if block == nil {
		return nil, errors.New("failed to decode PEM block containing public key")
	}

	// 检查PEM块类型
	switch block.Type {
	case "PUBLIC KEY": // PKIX/SPKI 格式 (Node.js生成的格式)
		pub, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		if rsaPub, ok := pub.(*rsa.PublicKey); ok {
			return rsaPub, nil
		}
		return nil, errors.New("not an RSA public key")

	case "RSA PUBLIC KEY": // PKCS#1 格式
		return x509.ParsePKCS1PublicKey(block.Bytes)

	default:
		// 尝试自动检测格式
		// 先尝试PKIX格式
		if pub, err := x509.ParsePKIXPublicKey(block.Bytes); err == nil {
			if rsaPub, ok := pub.(*rsa.PublicKey); ok {
				return rsaPub, nil
			}
		}

		// 再尝试PKCS#1格式
		if pub, err := x509.ParsePKCS1PublicKey(block.Bytes); err == nil {
			return pub, nil
		}

		return nil, errors.New("unsupported public key format: " + block.Type)
	}
}

func (pub *rsaPub) Encrypt(encryptStr string) (string, error) {
	publicKey, err := parsePublicKey([]byte(pub.PublicKey))
	if err != nil {
		return "", err
	}

	// 对明文进行加密
	encryptedStr, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, []byte(encryptStr))
	if err != nil {
		return "", err
	}

	// 返回密文
	return base64.StdEncoding.EncodeToString(encryptedStr), nil
}

func (pub *rsaPub) EncryptURLEncoding(encryptStr string) (string, error) {
	publicKey, err := parsePublicKey([]byte(pub.PublicKey))
	if err != nil {
		return "", err
	}

	// 对明文进行加密
	encryptedStr, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, []byte(encryptStr))
	if err != nil {
		return "", err
	}

	// 返回密文
	return base64.URLEncoding.EncodeToString(encryptedStr), nil
}

func (pri *rsaPri) i() {}

// parsePrivateKey 解析多种格式的私钥
func parsePrivateKey(pemData []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(pemData)
	if block == nil {
		return nil, errors.New("failed to decode PEM block containing private key")
	}

	// 检查PEM块类型
	switch block.Type {
	case "RSA PRIVATE KEY": // PKCS#1 格式
		return x509.ParsePKCS1PrivateKey(block.Bytes)

	case "PRIVATE KEY": // PKCS#8 格式 (Node.js生成的格式)
		priv, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		if rsaPriv, ok := priv.(*rsa.PrivateKey); ok {
			return rsaPriv, nil
		}
		return nil, errors.New("not an RSA private key")

	case "ENCRYPTED PRIVATE KEY": // 加密的PKCS#8格式
		// 注意：对于加密的PKCS#8私钥，需要先解密
		// 这里返回错误，建议用户先解密再使用
		return nil, errors.New("encrypted private key detected. Please decrypt it first before using")

	default:
		// 尝试自动检测格式
		// 先尝试PKCS#1格式
		if priv, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
			return priv, nil
		}

		// 再尝试PKCS#8格式
		if priv, err := x509.ParsePKCS8PrivateKey(block.Bytes); err == nil {
			if rsaPriv, ok := priv.(*rsa.PrivateKey); ok {
				return rsaPriv, nil
			}
			return nil, errors.New("not an RSA private key")
		}

		return nil, errors.New("unsupported private key format: " + block.Type)
	}
}

func (pri *rsaPri) Decrypt(decryptStr string) (string, error) {
	return pri.decryptInternal(decryptStr, base64.StdEncoding.DecodeString)
}

func (pri *rsaPri) DecryptURLEncoding(decryptStr string) (string, error) {
	return pri.decryptInternal(decryptStr, base64.URLEncoding.DecodeString)
}

// decryptInternal 内部解密方法
func (pri *rsaPri) decryptInternal(
	decryptStr string,
	decodeFunc func(string) ([]byte, error),
) (string, error) {
	// 解析私钥
	privateKey, err := parsePrivateKey([]byte(pri.PrivateKey))
	if err != nil {
		return "", err
	}

	// 解码Base64密文
	decryptBytes, err := decodeFunc(decryptStr)
	if err != nil {
		return "", err
	}

	// 对密文进行解密
	decrypted, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, decryptBytes)
	if err != nil {
		return "", err
	}

	// 返回明文
	return string(decrypted), nil
}

// DetectKeyFormat 辅助函数：检测密钥格式并提供使用建议
func DetectKeyFormat(key string) (string, string, error) {
	block, _ := pem.Decode([]byte(key))
	if block == nil {
		return "", "", errors.New("invalid PEM format")
	}

	var format string
	var suggestion string

	switch block.Type {
	case "PUBLIC KEY":
		if _, err := x509.ParsePKIXPublicKey(block.Bytes); err == nil {
			format = "PKIX/SPKI Public Key"
			suggestion = "可以直接使用"
		} else {
			format = "Unknown Public Key"
		}

	case "RSA PUBLIC KEY":
		format = "PKCS#1 Public Key"
		suggestion = "可以直接使用"

	case "RSA PRIVATE KEY":
		format = "PKCS#1 Private Key"
		suggestion = "可以直接使用"

	case "PRIVATE KEY":
		// 尝试解析以确认是否是PKCS#8
		if _, err := x509.ParsePKCS8PrivateKey(block.Bytes); err == nil {
			format = "PKCS#8 Private Key"
			suggestion = "可以直接使用（兼容Node.js生成的格式）"
		} else {
			format = "Unknown PKCS#8 Private Key"
		}

	case "ENCRYPTED PRIVATE KEY":
		format = "Encrypted PKCS#8 Private Key"
		suggestion = "请先使用OpenSSL解密：openssl pkcs8 -in encrypted.pem -out decrypted.pem -nocrypt"

	default:
		format = "Unknown Key Type: " + block.Type
		suggestion = "请转换为支持的格式"
	}

	return format, suggestion, nil
}

// ConvertPKCS8ToPKCS1 转换函数：PKCS#8 转 PKCS#1（如果需要）
func ConvertPKCS8ToPKCS1(pkcs8Key string) (string, error) {
	block, _ := pem.Decode([]byte(pkcs8Key))
	if block == nil {
		return "", errors.New("failed to decode PEM block")
	}

	if block.Type != "PRIVATE KEY" {
		return "", errors.New("not a PKCS#8 private key")
	}

	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return "", errors.New("not an RSA private key")
	}

	// 编码为 PKCS#1
	pkcs1Bytes := x509.MarshalPKCS1PrivateKey(rsaKey)
	pkcs1Block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: pkcs1Bytes,
	}

	return string(pem.EncodeToMemory(pkcs1Block)), nil
}

// IsEncryptedKey 检查是否是加密私钥
func IsEncryptedKey(key string) bool {
	block, _ := pem.Decode([]byte(key))
	if block == nil {
		return false
	}

	// 现代加密私钥使用 ENCRYPTED PRIVATE KEY 类型
	return block.Type == "ENCRYPTED PRIVATE KEY"
}

// DecryptPrivateKey 解密加密私钥的辅助函数（使用OpenSSL命令）
func DecryptPrivateKey(encryptedKey, passphrase string) (string, error) {
	if !IsEncryptedKey(encryptedKey) {
		return "", errors.New("not an encrypted private key")
	}

	// 注意：这里只提供建议，实际解密应在外部进行
	return "", errors.New(`请使用OpenSSL解密：
openssl pkcs8 -in encrypted.pem -out decrypted.pem -passin pass:"` + passphrase + `"

或者使用编程方式在外部解密后再传入本库`)
}