package edgeencrypt

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
)

const (
	KeySize = 256
)

// EncryptData encrypts data using RSA public key, using 256 bit AES key
func EncryptData(pemPublicKey string, data []byte) ([]byte, error) {
	rsaPublicKey, err := toRSAPublicKey(pemPublicKey)
	if err != nil {
		return nil, err
	}
	aesKey, err := geneRandomAES256BitKey()
	if err != nil {
		return nil, err
	}

	cryptoData, err := encryptData(data, aesKey)
	if err != nil {
		return nil, err
	}
	cryptoKey, err := encryptEASKey(rsaPublicKey, aesKey)
	if err != nil {
		return nil, err
	}

	encryptedData := append(cryptoKey, cryptoData...)
	return encryptedData, nil
}

func toRSAPublicKey(pemPublicKey string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pemPublicKey))
	pk, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %w", err)
	}
	rsaKey, ok := pk.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("public key is not in rsa format")
	}
	return rsaKey, nil
}

func geneRandomAES256BitKey() ([]byte, error) {
	key := make([]byte, 32) // AES-256
	if _, err := rand.Read(key); err != nil {
		return nil, fmt.Errorf("failed to generate random key: %w", err)
	}
	return key, nil
}

func encryptData(data, aesKey []byte) ([]byte, error) {
	blockCipher, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(blockCipher)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}
	return gcm.Seal(nonce, nonce, data, nil), nil
}

func encryptEASKey(rsaKey *rsa.PublicKey, aesKey []byte) ([]byte, error) {
	return rsa.EncryptOAEP(sha256.New(), rand.Reader, rsaKey, aesKey, nil)
}