package backup

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"

	"golang.org/x/crypto/argon2"
)

const (
	// Argon2id parameters — byte-identical to the cms settings_crypto so that
	// archives produced here decrypt through the existing cms restore path.
	argon2Time    = 3
	argon2Memory  = 64 * 1024 // 64 MB
	argon2Threads = 4
	argon2KeyLen  = 32 // AES-256
	saltLength    = 16
)

// EncryptionParams stores the parameters used for key derivation. Its JSON shape
// is wire-identical to the cms EncryptionParams so manifests interoperate.
type EncryptionParams struct {
	Algorithm   string `json:"algorithm"`
	KDF         string `json:"kdf"`
	Salt        string `json:"salt"`
	Iterations  int    `json:"iterations"`
	Memory      int    `json:"memory"`
	Parallelism int    `json:"parallelism"`
}

// deriveKey derives a 32-byte key from password using Argon2id.
func deriveKey(password string, salt []byte) []byte {
	return argon2.IDKey([]byte(password), salt, argon2Time, argon2Memory, argon2Threads, argon2KeyLen)
}

// EncryptWithPassword encrypts data with AES-256-GCM using a password-derived key.
// The nonce is prepended to the returned ciphertext.
func EncryptWithPassword(plaintext []byte, password string) ([]byte, *EncryptionParams, error) {
	// Generate random salt
	salt := make([]byte, saltLength)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, nil, fmt.Errorf("failed to generate salt: %w", err)
	}

	// Derive key from password
	key := deriveKey(password, salt)

	// Create AES cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	// Create GCM mode
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	// Generate nonce
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Encrypt (prepend nonce to ciphertext)
	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)

	return ciphertext, &EncryptionParams{
		Algorithm:   "aes-256-gcm",
		KDF:         "argon2id",
		Salt:        base64.StdEncoding.EncodeToString(salt),
		Iterations:  argon2Time,
		Memory:      argon2Memory,
		Parallelism: argon2Threads,
	}, nil
}

// DecryptWithPassword decrypts data encrypted with EncryptWithPassword.
func DecryptWithPassword(ciphertext []byte, password string, params *EncryptionParams) ([]byte, error) {
	// Decode salt
	salt, err := base64.StdEncoding.DecodeString(params.Salt)
	if err != nil {
		return nil, fmt.Errorf("failed to decode salt: %w", err)
	}

	// Derive key from password
	key := deriveKey(password, salt)

	// Create AES cipher
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher: %w", err)
	}

	// Create GCM mode
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	// Extract nonce
	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}
	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]

	// Decrypt
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("decryption failed - invalid password or corrupted data")
	}

	return plaintext, nil
}