package clusteragent

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

// EncryptWithPublicKey encrypts plaintext using an RSA public key in PEM format.
// The encryption uses OAEP with SHA-256, matching the cluster-agent's decryption.
func EncryptWithPublicKey(pubPEM []byte, plaintext []byte) ([]byte, error) {
	pubKey, err := parsePublicKeyFromPEM(pubPEM)
	if err != nil {
		return nil, err
	}
	ciphertext, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, plaintext, nil)
	if err != nil {
		return nil, fmt.Errorf("encryption failed: %w", err)
	}
	return ciphertext, nil
}

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")
	}

	var pkixErr error
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err == nil {
		pubKey, ok := pubInterface.(*rsa.PublicKey)
		if !ok {
			return nil, errors.New("parsed key is not an RSA public key")
		}
		return pubKey, nil
	}
	pkixErr = err

	pkcs1Key, err := x509.ParsePKCS1PublicKey(block.Bytes)
	if err == nil {
		return pkcs1Key, nil
	}

	return nil, fmt.Errorf("failed to parse RSA public key: pkix=%v pkcs1=%v", pkixErr, err)
}