// Package crypto provides API key encryption primitives for Switchboard Core.
//
// Two-tier model:
//
//   - Global/Team keys: AES-256-GCM with a key derived from ENCRYPTION_KEY env var.
//   - Personal (BYOK) keys: AES-256-GCM with a per-user encryption key (UEK).
//     The UEK is itself encrypted with a key derived from the user's password
//     via Argon2id.
//
// The UEK lives in server memory for the session duration and is never
// persisted in plaintext.
package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"

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

// Argon2id parameters — tuned for server-side derivation.
// Memory: 64 MiB, Iterations: 3, Parallelism: 4, Key length: 32 bytes.
const (
	argonMemory  = 64 * 1024 // 64 MiB
	argonTime    = 3
	argonThreads = 4
	argonKeyLen  = 32
	saltLen      = 16
	nonceLen     = 12 // AES-256-GCM standard nonce size
	uekLen       = 32 // 256-bit UEK
)

var (
	ErrDecryptionFailed = errors.New("decryption failed: wrong key or corrupted data")
	ErrVaultLocked      = errors.New("vault locked: user encryption key not in session")
	ErrNoEncryptionKey  = errors.New("ENCRYPTION_KEY not set: cannot encrypt/decrypt platform keys")
)

// --- Key Derivation ---

// DeriveKeyFromPassword derives a 256-bit key from a password and salt using Argon2id.
// Used for wrapping/unwrapping the per-user UEK.
func DeriveKeyFromPassword(password string, salt []byte) []byte {
	return argon2.IDKey([]byte(password), salt, argonTime, argonMemory, argonThreads, argonKeyLen)
}

// DeriveKeyFromEnv derives a 256-bit key from an environment variable value
// using HKDF-SHA256 with domain separation. Deterministic — same input always
// produces the same key.
func DeriveKeyFromEnv(envValue string) ([]byte, error) {
	if envValue == "" {
		return nil, ErrNoEncryptionKey
	}
	// HKDF with SHA-256: extract + expand with context string for domain separation.
	// No salt (nil) — the env value itself is the input keying material.
	hkdfReader := hkdf.New(sha256.New, []byte(envValue), nil, []byte("switchboard-api-key-encryption-v1"))
	key := make([]byte, 32)
	if _, err := io.ReadFull(hkdfReader, key); err != nil {
		return nil, fmt.Errorf("HKDF derivation failed: %w", err)
	}
	return key, nil
}

// --- Salt / Nonce Generation ---

// GenerateSalt returns a cryptographically random salt for Argon2id.
func GenerateSalt() ([]byte, error) {
	salt := make([]byte, saltLen)
	if _, err := rand.Read(salt); err != nil {
		return nil, fmt.Errorf("failed to generate salt: %w", err)
	}
	return salt, nil
}

// --- UEK Lifecycle ---

// GenerateUEK creates a new random 256-bit User Encryption Key.
func GenerateUEK() ([]byte, error) {
	uek := make([]byte, uekLen)
	if _, err := rand.Read(uek); err != nil {
		return nil, fmt.Errorf("failed to generate UEK: %w", err)
	}
	return uek, nil
}

// WrapUEK encrypts a UEK with a password-derived key (PDK).
// Returns (ciphertext, nonce, error).
func WrapUEK(uek, pdk []byte) ([]byte, []byte, error) {
	return encrypt(uek, pdk)
}

// UnwrapUEK decrypts a UEK using the password-derived key.
func UnwrapUEK(ciphertext, nonce, pdk []byte) ([]byte, error) {
	return decrypt(ciphertext, nonce, pdk)
}

// --- API Key Encryption ---

// Encrypt encrypts plaintext with the given 256-bit key using AES-256-GCM.
// Returns (ciphertext, nonce, error).
func Encrypt(plaintext string, key []byte) ([]byte, []byte, error) {
	if plaintext == "" {
		return nil, nil, nil // Nothing to encrypt
	}
	return encrypt([]byte(plaintext), key)
}

// Decrypt decrypts ciphertext with the given 256-bit key and nonce.
func Decrypt(ciphertext, nonce, key []byte) (string, error) {
	if len(ciphertext) == 0 {
		return "", nil // Nothing to decrypt
	}
	plain, err := decrypt(ciphertext, nonce, key)
	if err != nil {
		return "", err
	}
	return string(plain), nil
}

// --- Internal AES-256-GCM ---

func encrypt(plaintext, key []byte) ([]byte, []byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, nil, fmt.Errorf("aes.NewCipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, nil, fmt.Errorf("cipher.NewGCM: %w", err)
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, nil, fmt.Errorf("nonce generation: %w", err)
	}

	ciphertext := gcm.Seal(nil, nonce, plaintext, nil)
	return ciphertext, nonce, nil
}

func decrypt(ciphertext, nonce, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("aes.NewCipher: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("cipher.NewGCM: %w", err)
	}

	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, ErrDecryptionFailed
	}

	return plaintext, nil
}