package service

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"

	"github.com/tidbcloud/serverless-service/internal/domain/vo"
)

const (
	exportCreatePayloadDataKeySize = 32
	exportCreatePayloadAAD         = "devtier-service/export-create-payload"
)

type exportCreatePayloadEnvelope struct {
	EncryptedDataKey string `json:"encrypted_data_key"`
	Nonce            string `json:"nonce"`
	Ciphertext       string `json:"ciphertext"`
}

// AWS KMS Encrypt accepts at most 4096 bytes, so only the random AES data key
// is sent to KMS. The export payload itself is authenticated and encrypted locally.
func (e *exportService) encryptExportCreatePayload(export *vo.Export) (string, error) {
	payload, err := json.Marshal(export)
	if err != nil {
		return "", fmt.Errorf("marshal export create payload: %w", err)
	}

	dataKey := make([]byte, exportCreatePayloadDataKeySize)
	if _, err := rand.Read(dataKey); err != nil {
		return "", fmt.Errorf("generate export create payload data key: %w", err)
	}

	gcm, err := newExportCreatePayloadGCM(dataKey)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", fmt.Errorf("generate export create payload nonce: %w", err)
	}
	ciphertext := gcm.Seal(nil, nonce, payload, []byte(exportCreatePayloadAAD))

	encryptedDataKey, err := e.encryptorFacade.EncryptBase64(string(dataKey))
	if err != nil {
		return "", fmt.Errorf("encrypt export create payload data key: %w", err)
	}

	envelope, err := json.Marshal(exportCreatePayloadEnvelope{
		EncryptedDataKey: encryptedDataKey,
		Nonce:            base64.StdEncoding.EncodeToString(nonce),
		Ciphertext:       base64.StdEncoding.EncodeToString(ciphertext),
	})
	if err != nil {
		return "", fmt.Errorf("marshal export create payload envelope: %w", err)
	}
	return string(envelope), nil
}

func (e *exportService) decryptExportCreatePayload(encryptedPayload string) (*vo.Export, error) {
	if encryptedPayload == "" {
		return nil, errors.New("empty export create payload")
	}

	var envelope exportCreatePayloadEnvelope
	if err := json.Unmarshal([]byte(encryptedPayload), &envelope); err != nil {
		return nil, fmt.Errorf("unmarshal export create payload envelope: %w", err)
	}
	if envelope.EncryptedDataKey == "" || envelope.Nonce == "" || envelope.Ciphertext == "" {
		return nil, errors.New("incomplete export create payload envelope")
	}

	dataKey, err := e.encryptorFacade.DecryptBase64(envelope.EncryptedDataKey)
	if err != nil {
		return nil, fmt.Errorf("decrypt export create payload data key: %w", err)
	}
	gcm, err := newExportCreatePayloadGCM([]byte(dataKey))
	if err != nil {
		return nil, err
	}

	nonce, err := base64.StdEncoding.DecodeString(envelope.Nonce)
	if err != nil {
		return nil, fmt.Errorf("decode export create payload nonce: %w", err)
	}
	if len(nonce) != gcm.NonceSize() {
		return nil, fmt.Errorf("invalid export create payload nonce size: %d", len(nonce))
	}
	ciphertext, err := base64.StdEncoding.DecodeString(envelope.Ciphertext)
	if err != nil {
		return nil, fmt.Errorf("decode export create payload ciphertext: %w", err)
	}
	payload, err := gcm.Open(nil, nonce, ciphertext, []byte(exportCreatePayloadAAD))
	if err != nil {
		return nil, fmt.Errorf("decrypt export create payload: %w", err)
	}

	export := &vo.Export{}
	if err := json.Unmarshal(payload, export); err != nil {
		return nil, fmt.Errorf("unmarshal export create payload: %w", err)
	}
	return export, nil
}

func newExportCreatePayloadGCM(dataKey []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(dataKey)
	if err != nil {
		return nil, fmt.Errorf("initialize export create payload cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("initialize export create payload GCM: %w", err)
	}
	return gcm, nil
}