package rsa

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"github.com/gogf/gf/v2/frame/g"
	"os"
	"strings"
	"sync"
)

type rsaClient struct {
	publicKey  *rsa.PublicKey
	privateKey *rsa.PrivateKey
}

var (
	instance *rsaClient
	once     sync.Once
)

// init 会在包初始化时自动调用，用于加载默认的 RSA 公钥和私钥。
func init() {
	ctx := context.Background()
	once.Do(func() {
		instance = &rsaClient{}
		err := instance.loadKeys(g.Config().MustGet(ctx, "gamelife.rsaKey").String())
		if err != nil {
			panic("加载 RSA 密钥失败: " + err.Error())
		}
	})
}

// GetRsaClient 返回 RSA 客户端的单例实例。
//
// 通常用于执行加解密操作。
func GetRsaClient() *rsaClient {
	return instance
}

// EncryptWithRsaPublicKey 使用加载的 RSA 公钥对原始数据进行加密。
//
// 参数：
//   - plain: 待加密的明文数据。
//
// 返回值：
//   - 加密后的密文数据。
//   - 如果加密失败，则返回错误。
func (c *rsaClient) EncryptWithRsaPublicKey(plain []byte) ([]byte, error) {
	if c.publicKey == nil {
		return nil, errors.New("公钥未加载")
	}
	return rsa.EncryptPKCS1v15(rand.Reader, c.publicKey, plain)
}

// DecryptWithRsaPrivateKey 使用加载的 RSA 私钥对密文数据进行解密。
//
// 参数：
//   - cipher: 加密后的密文数据。
//
// 返回值：
//   - 解密后的明文数据。
//   - 如果解密失败，则返回错误。
func (c *rsaClient) DecryptWithRsaPrivateKey(cipher []byte) ([]byte, error) {
	if c.privateKey == nil {
		return nil, errors.New("私钥未加载")
	}
	return rsa.DecryptPKCS1v15(rand.Reader, c.privateKey, cipher)
}

// loadKeys 从指定文件中加载 RSA 公钥和私钥。
//
// 参数：
//   - publicKeyPath: 公钥 PEM 文件路径。
//   - privateKeyPath: 私钥 PEM 文件路径。
//
// 返回值：
//   - 成功返回 nil，否则返回错误信息。
func (c *rsaClient) loadKeys(keyFilePath string) error {
	// 读取密钥文件
	keyBytes, err := os.ReadFile(keyFilePath)
	if err != nil {
		return fmt.Errorf("读取密钥文件失败: %w", err)
	}

	var pubFound, privFound bool
	rest := keyBytes

	for {
		var block *pem.Block
		block, rest = pem.Decode(rest)
		if block == nil {
			break
		}

		// 转换为小写以匹配你的密钥格式
		blockTypeLower := strings.ToLower(block.Type)

		switch blockTypeLower {
		case "rsa public key":
			// 尝试解析 PKCS#1 格式的公钥
			pubKey, err := x509.ParsePKCS1PublicKey(block.Bytes)
			if err != nil {
				// 尝试 PKIX 格式（兼容性处理）
				pub, err := x509.ParsePKIXPublicKey(block.Bytes)
				if err != nil {
					return fmt.Errorf("无法解析 RSA 公钥: %w", err)
				}
				var ok bool
				if c.publicKey, ok = pub.(*rsa.PublicKey); !ok {
					return errors.New("解析的公钥不是 RSA 公钥")
				}
			} else {
				c.publicKey = pubKey
			}
			pubFound = true

		case "rsa private key":
			// 解析 PKCS#1 格式的私钥
			privKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
			if err != nil {
				return fmt.Errorf("无法解析 RSA 私钥: %w", err)
			}
			c.privateKey = privKey
			privFound = true

		default:
			// 忽略未知的 PEM 块
			continue
		}
	}

	if !pubFound {
		return errors.New("未找到有效的 RSA 公钥")
	}
	if !privFound {
		return errors.New("未找到有效的 RSA 私钥")
	}

	return nil
}

// EncryptWithPublicKey 使用公钥加密数据
func (c *rsaClient) EncryptWithPublicKey(plaintext []byte) ([]byte, error) {
	if c.publicKey == nil {
		return nil, errors.New("公钥未加载")
	}

	// 使用 OAEP 填充进行加密
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, c.publicKey, plaintext)
	if err != nil {
		return nil, fmt.Errorf("公钥加密失败: %w", err)
	}
	return ciphertext, nil
}

// DecryptWithPrivateKey 使用私钥解密数据
func (c *rsaClient) DecryptWithPrivateKey(ciphertext []byte) ([]byte, error) {
	if c.privateKey == nil {
		return nil, errors.New("私钥未加载")
	}

	// 使用 OAEP 填充进行解密
	plaintext, err := rsa.DecryptPKCS1v15(rand.Reader, c.privateKey, ciphertext)
	if err != nil {
		return nil, fmt.Errorf("私钥解密失败: %w", err)
	}
	return plaintext, nil
}