package crypto

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

const (
	defaultKeyBits = 2048
)

type Encryptor struct {
	privateKey *rsa.PrivateKey
}

func NewEncryptor() (*Encryptor, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, defaultKeyBits)
	if err != nil {
		return nil, err
	}
	return &Encryptor{
		privateKey: privateKey,
	}, nil
}

func (e *Encryptor) ExportPublicKeyAsPEM() ([]byte, error) {
	if e.privateKey == nil {
		return nil, errors.New("private key is not initialized")
	}

	pubASN1, err := x509.MarshalPKIXPublicKey(&e.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
}

func (e *Encryptor) Decrypt(encryptedData []byte) ([]byte, error) {
	if e.privateKey == nil {
		return nil, errors.New("private key not initialized, cannot decrypt")
	}

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

// parsePublicKeyFromPEM parses an RSA public key from a PEM-encoded byte slice.
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
}