package edgeencrypt

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"fmt"
)

// DecryptFunc is a function that decrypts data, key decryption is handled externally by GCP
type DecryptFunc func(ctx context.Context, keyRing, key, keyVersion string, data []byte) ([]byte, error)

// DecryptData data using RSA private key, using 256 bit AES key
func DecryptData(ctx context.Context, e *EncryptedData, ec *EncryptionClaims, decrypt DecryptFunc) ([]byte, error) {
	if ec.Role != Decryption {
		return nil, fmt.Errorf("invalid role for decryption: %s", ec.Role)
	}
	if len(e.Data) < KeySize {
		return nil, fmt.Errorf("encrypted data is less than key size")
	}
	encryptedKey, encryptedData := splitEncryptedData(e.Data)
	decryptedKey, err := decrypt(ctx, e.BannerEdgeID, e.ChannelID, e.KeyVersion, encryptedKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt key: %w", err)
	}
	decryptedData, err := decryptData(encryptedData, decryptedKey)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt data: %w", err)
	}
	return decryptedData, nil
}

func decryptData(data, aesKey []byte) ([]byte, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher block: %w", err)
	}
	nonceSize := gcm.NonceSize()
	if len(data) < nonceSize {
		return nil, fmt.Errorf("data is too short")
	}
	nonce, data := data[:nonceSize], data[nonceSize:]
	plaintext, err := gcm.Open(nil, nonce, data, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt data: %w", err)
	}
	return plaintext, nil
}

// splitEncryptedData splits data into two parts, first part is 256 bit AES key, second part is encrypted data
func splitEncryptedData(data []byte) ([]byte, []byte) {
	return data[:KeySize], data[KeySize:]
}