package crypto

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

/*
  It's copied from internal/crypto/.
*/

const PEMUpdatedErrorMsg = "PEMUpdated"

var PEMUpdatedError = errors.New(PEMUpdatedErrorMsg)

// RSAEncryptor encapsulates the RSA private key for server-side operations.
// It is designed to be a temporary key manager residing in memory.
type RSAEncryptor struct {
	privateKey *rsa.PrivateKey
}

// NewRSAEncryptor generates a new RSA key pair of the specified bit size and returns
// an RSAEncryptor instance holding the private key.
//
// The bitSize parameter determines the key strength, with 2048 or 4096 being recommended for security.
func NewRSAEncryptor(bitSize int) (*RSAEncryptor, error) {
	if bitSize < 2048 {
		return nil, errors.New("key bit size must be at least 2048 for security")
	}

	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate RSA key pair: %w", err)
	}

	return &RSAEncryptor{
		privateKey: privateKey,
	}, nil
}

// ExportPublicKeyAsPEM exports the RSA public key to a PEM-encoded byte slice.
//
// This method is typically used by the server to provide its public key to clients.
func (r *RSAEncryptor) ExportPublicKeyAsPEM() ([]byte, error) {
	if r.privateKey == nil {
		return nil, errors.New("private key is not initialized")
	}

	pubASN1, err := x509.MarshalPKIXPublicKey(&r.privateKey.PublicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal public key to ASN.1: %w", err)
	}

	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubASN1,
	})

	return pubPEM, nil
}

// Decrypt uses the encapsulated private key to decrypt the provided data.
//
// This method is exclusively for server-side use.
func (r *RSAEncryptor) Decrypt(encryptedData []byte) ([]byte, error) {
	if r.privateKey == nil {
		return nil, errors.New("private key not initialized, cannot decrypt")
	}

	decryptedData, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, r.privateKey, encryptedData, nil)
	if err != nil {
		return nil, fmt.Errorf("decryption failed: %w", err)
	}
	return decryptedData, nil
}

// VerifyKeyMatch checks if the provided public key PEM matches the private key held by the encryptor.
//
// It performs this check by signing a fixed piece of data with the private key
// and then verifying the signature with the public key. This method is crucial
// for ensuring a client is using the correct public key for the server's current key pair.
func (r *RSAEncryptor) VerifyKeyMatch(pubPEM []byte) (bool, error) {
	if r.privateKey == nil {
		return false, errors.New("private key is not initialized")
	}

	parsedPubKey, err := ParsePublicKeyFromPEM(pubPEM)
	if err != nil {
		return false, fmt.Errorf("failed to parse public key PEM: %w", err)
	}

	dataToSign := []byte("test_data_for_key_match_verification")
	hashed := sha256.Sum256(dataToSign)
	signature, err := rsa.SignPSS(rand.Reader, r.privateKey, crypto.SHA256, hashed[:], nil)
	if err != nil {
		return false, fmt.Errorf("failed to sign with private key: %w", err)
	}

	err = rsa.VerifyPSS(parsedPubKey, crypto.SHA256, hashed[:], signature, nil)
	if err != nil {
		return false, nil
	}
	return true, nil
}

// ParsePublicKeyFromPEM parses an RSA public key from a PEM-encoded byte slice.
// This is an internal helper function and is not typically called directly by the client.
func ParsePublicKeyFromPEM(pubPEM []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pubPEM)
	if block == nil {
		return nil, errors.New("PEM data decoding failed, possibly due to incorrect format")
	}
	if block.Type != "RSA PUBLIC KEY" {
		return nil, fmt.Errorf("incorrect PEM block type: %s", block.Type)
	}

	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %w", err)
	}

	pubKey, ok := pubInterface.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("parsed key is not an RSA public key")
	}
	return pubKey, nil
}

// EncryptWithPEM encrypts the given data using a PEM-encoded RSA public key.
//
// This is a convenient function for clients, as it handles the PEM parsing internally.
// The client only needs to provide the public key in its raw PEM format.
func EncryptWithPEM(pubPEM []byte, data []byte) ([]byte, error) {
	publicKey, err := ParsePublicKeyFromPEM(pubPEM)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key from PEM: %w", err)
	}

	encryptedData, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, publicKey, data, nil)
	if err != nil {
		return nil, fmt.Errorf("encryption failed: %w", err)
	}
	return encryptedData, nil
}