// Package user implements the user-to-user encryption engine.
package user

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/hkdf"

	"git.wntrmute.dev/kyle/metacrypt/internal/barrier"
	"git.wntrmute.dev/kyle/metacrypt/internal/crypto"
	"git.wntrmute.dev/kyle/metacrypt/internal/engine"
)

const (
	maxRecipients = 100
	nonceSize     = 12
	keySize       = 32
	hkdfInfoPrefix = "metacrypt-user-v1:"
)

var (
	ErrSealed         = errors.New("user: engine is sealed")
	ErrForbidden      = errors.New("user: forbidden")
	ErrUnauthorized   = errors.New("user: authentication required")
	ErrUserNotFound   = errors.New("user: user not found")
	ErrUserExists     = errors.New("user: user already exists")
	ErrTooMany        = errors.New("user: too many recipients")
	ErrInvalidEnvelope = errors.New("user: invalid envelope")
	ErrRecipientNotFound = errors.New("user: recipient entry not found in envelope")
	ErrDecryptionFailed  = errors.New("user: decryption failed")
	ErrInvalidAlgorithm  = errors.New("user: unsupported algorithm")
	ErrNoRecipients      = errors.New("user: no recipients specified")
)

// userState holds in-memory state for a loaded user.
type userState struct {
	privKey *ecdh.PrivateKey
	pubKey  *ecdh.PublicKey
	config  *UserKeyConfig
}

// UserEngine implements the user-to-user encryption engine.
type UserEngine struct {
	barrier   barrier.Barrier
	config    *UserConfig
	users     map[string]*userState
	mountPath string
	mu        sync.RWMutex
}

// NewUserEngine creates a new user engine instance.
func NewUserEngine() engine.Engine {
	return &UserEngine{
		users: make(map[string]*userState),
	}
}

func (e *UserEngine) Type() engine.EngineType {
	return engine.EngineTypeUser
}

func (e *UserEngine) Initialize(ctx context.Context, b barrier.Barrier, mountPath string, config map[string]interface{}) error {
	cfg := &UserConfig{
		KeyAlgorithm: "x25519",
		SymAlgorithm: "aes256-gcm",
	}
	if v, ok := config["key_algorithm"].(string); ok && v != "" {
		cfg.KeyAlgorithm = v
	}
	if v, ok := config["sym_algorithm"].(string); ok && v != "" {
		cfg.SymAlgorithm = v
	}

	if err := validateKeyAlgorithm(cfg.KeyAlgorithm); err != nil {
		return err
	}
	if cfg.SymAlgorithm != "aes256-gcm" {
		return fmt.Errorf("user: unsupported symmetric algorithm: %s", cfg.SymAlgorithm)
	}

	data, err := json.Marshal(cfg)
	if err != nil {
		return fmt.Errorf("user: marshal config: %w", err)
	}
	if err := b.Put(ctx, mountPath+"config.json", data); err != nil {
		return fmt.Errorf("user: store config: %w", err)
	}

	e.barrier = b
	e.config = cfg
	e.mountPath = mountPath
	e.users = make(map[string]*userState)
	return nil
}

func (e *UserEngine) Unseal(ctx context.Context, b barrier.Barrier, mountPath string) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	e.barrier = b
	e.mountPath = mountPath

	// Load config.
	data, err := b.Get(ctx, mountPath+"config.json")
	if err != nil {
		return fmt.Errorf("user: load config: %w", err)
	}
	var cfg UserConfig
	if err := json.Unmarshal(data, &cfg); err != nil {
		return fmt.Errorf("user: parse config: %w", err)
	}
	e.config = &cfg
	e.users = make(map[string]*userState)

	// Load all user keys.
	prefix := mountPath + "users/"
	paths, err := b.List(ctx, prefix)
	if err != nil {
		return nil // no users yet
	}

	// Discover unique usernames from paths like "alice/config.json", "alice/priv.key".
	seen := make(map[string]bool)
	for _, p := range paths {
		parts := strings.SplitN(p, "/", 2)
		if len(parts) > 0 && parts[0] != "" {
			seen[parts[0]] = true
		}
	}

	for username := range seen {
		if err := e.loadUser(ctx, username); err != nil {
			return fmt.Errorf("user: load user %q: %w", username, err)
		}
	}

	return nil
}

func (e *UserEngine) Seal() error {
	e.mu.Lock()
	defer e.mu.Unlock()

	// Zeroize all private keys.
	for _, u := range e.users {
		if u.privKey != nil {
			raw := u.privKey.Bytes()
			crypto.Zeroize(raw)
		}
	}
	e.users = nil
	e.config = nil
	e.barrier = nil
	return nil
}

func (e *UserEngine) HandleRequest(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	switch req.Operation {
	case "register":
		return e.handleRegister(ctx, req)
	case "provision":
		return e.handleProvision(ctx, req)
	case "get-public-key":
		return e.handleGetPublicKey(ctx, req)
	case "list-users":
		return e.handleListUsers(ctx, req)
	case "encrypt":
		return e.handleEncrypt(ctx, req)
	case "decrypt":
		return e.handleDecrypt(ctx, req)
	case "re-encrypt":
		return e.handleReEncrypt(ctx, req)
	case "rotate-key":
		return e.handleRotateKey(ctx, req)
	case "delete-user":
		return e.handleDeleteUser(ctx, req)
	default:
		return nil, fmt.Errorf("user: unknown operation: %s", req.Operation)
	}
}

// --- Operation handlers ---

func (e *UserEngine) handleRegister(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}

	username := req.CallerInfo.Username
	e.mu.RLock()
	if u, ok := e.users[username]; ok {
		pubB64 := base64.StdEncoding.EncodeToString(u.pubKey.Bytes())
		e.mu.RUnlock()
		return &engine.Response{Data: map[string]interface{}{
			"username":   username,
			"public_key": pubB64,
			"algorithm":  u.config.Algorithm,
		}}, nil
	}
	e.mu.RUnlock()

	e.mu.Lock()
	defer e.mu.Unlock()

	// Double-check after acquiring write lock.
	if u, ok := e.users[username]; ok {
		pubB64 := base64.StdEncoding.EncodeToString(u.pubKey.Bytes())
		return &engine.Response{Data: map[string]interface{}{
			"username":   username,
			"public_key": pubB64,
			"algorithm":  u.config.Algorithm,
		}}, nil
	}

	u, err := e.createUser(ctx, username, false)
	if err != nil {
		return nil, err
	}
	pubB64 := base64.StdEncoding.EncodeToString(u.pubKey.Bytes())
	return &engine.Response{Data: map[string]interface{}{
		"username":   username,
		"public_key": pubB64,
		"algorithm":  u.config.Algorithm,
	}}, nil
}

func (e *UserEngine) handleProvision(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}
	if !req.CallerInfo.IsAdmin {
		return nil, ErrForbidden
	}

	username, _ := req.Data["username"].(string)
	if username == "" {
		return nil, fmt.Errorf("user: username is required")
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	// No-op if exists.
	if u, ok := e.users[username]; ok {
		pubB64 := base64.StdEncoding.EncodeToString(u.pubKey.Bytes())
		return &engine.Response{Data: map[string]interface{}{
			"username":   username,
			"public_key": pubB64,
			"algorithm":  u.config.Algorithm,
		}}, nil
	}

	u, err := e.createUser(ctx, username, false)
	if err != nil {
		return nil, err
	}
	pubB64 := base64.StdEncoding.EncodeToString(u.pubKey.Bytes())
	return &engine.Response{Data: map[string]interface{}{
		"username":   username,
		"public_key": pubB64,
		"algorithm":  u.config.Algorithm,
	}}, nil
}

func (e *UserEngine) handleGetPublicKey(_ context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}
	if !req.CallerInfo.IsUser() {
		return nil, ErrForbidden
	}

	username, _ := req.Data["username"].(string)
	if username == "" {
		return nil, fmt.Errorf("user: username is required")
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	u, ok := e.users[username]
	if !ok {
		return nil, ErrUserNotFound
	}

	pubB64 := base64.StdEncoding.EncodeToString(u.pubKey.Bytes())
	return &engine.Response{Data: map[string]interface{}{
		"username":   username,
		"public_key": pubB64,
		"algorithm":  u.config.Algorithm,
	}}, nil
}

func (e *UserEngine) handleListUsers(_ context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}
	if !req.CallerInfo.IsUser() {
		return nil, ErrForbidden
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	usernames := make([]interface{}, 0, len(e.users))
	for name := range e.users {
		usernames = append(usernames, name)
	}

	return &engine.Response{Data: map[string]interface{}{
		"users": usernames,
	}}, nil
}

func (e *UserEngine) handleEncrypt(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}
	if !req.CallerInfo.IsUser() {
		return nil, ErrForbidden
	}

	plaintext, _ := req.Data["plaintext"].(string)
	if plaintext == "" {
		return nil, fmt.Errorf("user: plaintext is required")
	}
	metadata, _ := req.Data["metadata"].(string)

	recipientNames, err := extractRecipients(req.Data)
	if err != nil {
		return nil, err
	}
	if len(recipientNames) == 0 {
		return nil, ErrNoRecipients
	}
	if len(recipientNames) > maxRecipients {
		return nil, ErrTooMany
	}

	sender := req.CallerInfo.Username

	// Policy check for each recipient.
	if req.CheckPolicy != nil {
		for _, r := range recipientNames {
			resource := fmt.Sprintf("user/%s/recipient/%s", e.mountPath, r)
			effect, matched := req.CheckPolicy(resource, "write")
			if matched && effect == "deny" {
				return nil, fmt.Errorf("user: forbidden: policy denies encryption to recipient %s", r)
			}
		}
	}

	e.mu.Lock()

	// Auto-provision sender if not registered.
	if _, ok := e.users[sender]; !ok {
		if _, err := e.createUser(ctx, sender, true); err != nil {
			e.mu.Unlock()
			return nil, fmt.Errorf("user: auto-provision sender: %w", err)
		}
	}

	// Auto-provision recipients without keys.
	for _, r := range recipientNames {
		if _, ok := e.users[r]; !ok {
			if _, err := e.createUser(ctx, r, true); err != nil {
				e.mu.Unlock()
				return nil, fmt.Errorf("user: auto-provision recipient %s: %w", r, err)
			}
		}
	}

	senderState := e.users[sender]
	recipientStates := make(map[string]*userState, len(recipientNames))
	for _, r := range recipientNames {
		recipientStates[r] = e.users[r]
	}
	e.mu.Unlock()

	// Generate random DEK.
	dek := make([]byte, keySize)
	if _, err := rand.Read(dek); err != nil {
		return nil, fmt.Errorf("user: generate DEK: %w", err)
	}
	defer crypto.Zeroize(dek)

	// Encrypt plaintext with DEK.
	var aad []byte
	if metadata != "" {
		aad = []byte(metadata)
	}
	ct, err := encryptAESGCM(dek, []byte(plaintext), aad)
	if err != nil {
		return nil, fmt.Errorf("user: encrypt plaintext: %w", err)
	}

	// Wrap DEK for each recipient.
	recipients := make(map[string]*recipientEntry, len(recipientNames))
	for _, rName := range recipientNames {
		rState := recipientStates[rName]
		entry, wrapErr := wrapDEKForRecipient(senderState.privKey, rState.pubKey, dek, sender, rName)
		if wrapErr != nil {
			return nil, fmt.Errorf("user: wrap DEK for %s: %w", rName, wrapErr)
		}
		recipients[rName] = entry
	}

	env := &envelope{
		Version:      1,
		Sender:       sender,
		SymAlgorithm: e.config.SymAlgorithm,
		Ciphertext:   base64.StdEncoding.EncodeToString(ct),
		Metadata:     metadata,
		Recipients:   recipients,
	}

	envJSON, err := json.Marshal(env)
	if err != nil {
		return nil, fmt.Errorf("user: marshal envelope: %w", err)
	}

	envB64 := base64.StdEncoding.EncodeToString(envJSON)
	return &engine.Response{Data: map[string]interface{}{
		"envelope": envB64,
	}}, nil
}

func (e *UserEngine) handleDecrypt(_ context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}

	caller := req.CallerInfo.Username

	envelopeB64, _ := req.Data["envelope"].(string)
	if envelopeB64 == "" {
		return nil, fmt.Errorf("user: envelope is required")
	}

	env, err := parseEnvelope(envelopeB64)
	if err != nil {
		return nil, err
	}

	// Self-only: caller must be a recipient.
	entry, ok := env.Recipients[caller]
	if !ok {
		return nil, ErrRecipientNotFound
	}

	e.mu.RLock()
	callerState, callerExists := e.users[caller]
	senderState, senderExists := e.users[env.Sender]
	e.mu.RUnlock()

	if !callerExists {
		return nil, ErrUserNotFound
	}
	if !senderExists {
		return nil, fmt.Errorf("user: sender %q not found", env.Sender)
	}

	// Unwrap DEK.
	dek, err := unwrapDEK(callerState.privKey, senderState.pubKey, entry, env.Sender, caller)
	if err != nil {
		return nil, fmt.Errorf("user: unwrap DEK: %w", err)
	}
	defer crypto.Zeroize(dek)

	// Decrypt ciphertext.
	ct, err := base64.StdEncoding.DecodeString(env.Ciphertext)
	if err != nil {
		return nil, ErrInvalidEnvelope
	}

	var aad []byte
	if env.Metadata != "" {
		aad = []byte(env.Metadata)
	}

	plaintext, err := decryptAESGCM(dek, ct, aad)
	if err != nil {
		return nil, ErrDecryptionFailed
	}

	resp := map[string]interface{}{
		"plaintext": string(plaintext),
		"sender":    env.Sender,
	}
	if env.Metadata != "" {
		resp["metadata"] = env.Metadata
	}
	return &engine.Response{Data: resp}, nil
}

func (e *UserEngine) handleReEncrypt(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}

	caller := req.CallerInfo.Username

	envelopeB64, _ := req.Data["envelope"].(string)
	if envelopeB64 == "" {
		return nil, fmt.Errorf("user: envelope is required")
	}

	env, err := parseEnvelope(envelopeB64)
	if err != nil {
		return nil, err
	}

	// Self-only: caller must be a recipient.
	entry, ok := env.Recipients[caller]
	if !ok {
		return nil, ErrRecipientNotFound
	}

	e.mu.RLock()
	callerState, callerExists := e.users[caller]
	senderState, senderExists := e.users[env.Sender]
	e.mu.RUnlock()

	if !callerExists {
		return nil, ErrUserNotFound
	}
	if !senderExists {
		return nil, fmt.Errorf("user: sender %q not found", env.Sender)
	}

	// Unwrap DEK using old keys.
	dek, err := unwrapDEK(callerState.privKey, senderState.pubKey, entry, env.Sender, caller)
	if err != nil {
		return nil, fmt.Errorf("user: unwrap DEK: %w", err)
	}
	defer crypto.Zeroize(dek)

	// Verify we can decrypt (validates envelope integrity).
	ct, err := base64.StdEncoding.DecodeString(env.Ciphertext)
	if err != nil {
		return nil, ErrInvalidEnvelope
	}
	var aad []byte
	if env.Metadata != "" {
		aad = []byte(env.Metadata)
	}
	plaintext, err := decryptAESGCM(dek, ct, aad)
	if err != nil {
		return nil, ErrDecryptionFailed
	}

	// Generate new DEK and re-encrypt.
	newDEK := make([]byte, keySize)
	if _, err := rand.Read(newDEK); err != nil {
		return nil, fmt.Errorf("user: generate new DEK: %w", err)
	}
	defer crypto.Zeroize(newDEK)

	newCT, err := encryptAESGCM(newDEK, plaintext, aad)
	if err != nil {
		return nil, fmt.Errorf("user: re-encrypt: %w", err)
	}

	// Re-wrap for same recipients with current keys, using caller as new sender.
	e.mu.RLock()
	newRecipients := make(map[string]*recipientEntry, len(env.Recipients))
	for rName := range env.Recipients {
		rState, exists := e.users[rName]
		if !exists {
			e.mu.RUnlock()
			return nil, fmt.Errorf("user: recipient %q not found for re-encrypt", rName)
		}
		re, wrapErr := wrapDEKForRecipient(callerState.privKey, rState.pubKey, newDEK, caller, rName)
		if wrapErr != nil {
			e.mu.RUnlock()
			return nil, fmt.Errorf("user: re-wrap DEK for %s: %w", rName, wrapErr)
		}
		newRecipients[rName] = re
	}
	e.mu.RUnlock()

	newEnv := &envelope{
		Version:      1,
		Sender:       caller,
		SymAlgorithm: env.SymAlgorithm,
		Ciphertext:   base64.StdEncoding.EncodeToString(newCT),
		Metadata:     env.Metadata,
		Recipients:   newRecipients,
	}

	envJSON, err := json.Marshal(newEnv)
	if err != nil {
		return nil, fmt.Errorf("user: marshal envelope: %w", err)
	}
	envB64 := base64.StdEncoding.EncodeToString(envJSON)
	return &engine.Response{Data: map[string]interface{}{
		"envelope": envB64,
	}}, nil
}

func (e *UserEngine) handleRotateKey(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}

	caller := req.CallerInfo.Username

	e.mu.Lock()
	defer e.mu.Unlock()

	oldState, ok := e.users[caller]
	if !ok {
		return nil, ErrUserNotFound
	}

	// Generate new keypair.
	priv, err := generateKey(e.config.KeyAlgorithm)
	if err != nil {
		return nil, fmt.Errorf("user: rotate key: %w", err)
	}

	// Store new keys in barrier.
	if err := e.storeUserKeys(ctx, caller, priv, oldState.config.AutoProvisioned); err != nil {
		return nil, fmt.Errorf("user: rotate key: %w", err)
	}

	// Zeroize old key.
	oldRaw := oldState.privKey.Bytes()
	crypto.Zeroize(oldRaw)

	// Update in-memory state.
	e.users[caller] = &userState{
		privKey: priv,
		pubKey:  priv.PublicKey(),
		config: &UserKeyConfig{
			Algorithm:       e.config.KeyAlgorithm,
			CreatedAt:       time.Now().UTC(),
			AutoProvisioned: oldState.config.AutoProvisioned,
		},
	}

	pubB64 := base64.StdEncoding.EncodeToString(priv.PublicKey().Bytes())
	return &engine.Response{Data: map[string]interface{}{
		"username":   caller,
		"public_key": pubB64,
		"algorithm":  e.config.KeyAlgorithm,
	}}, nil
}

func (e *UserEngine) handleDeleteUser(ctx context.Context, req *engine.Request) (*engine.Response, error) {
	if req.CallerInfo == nil {
		return nil, ErrUnauthorized
	}
	if !req.CallerInfo.IsAdmin {
		return nil, ErrForbidden
	}

	username, _ := req.Data["username"].(string)
	if username == "" {
		return nil, fmt.Errorf("user: username is required")
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	oldState, ok := e.users[username]
	if !ok {
		return nil, ErrUserNotFound
	}

	// Zeroize private key.
	oldRaw := oldState.privKey.Bytes()
	crypto.Zeroize(oldRaw)

	// Delete from barrier.
	prefix := e.mountPath + "users/" + username + "/"
	paths, err := e.barrier.List(ctx, prefix)
	if err == nil {
		for _, p := range paths {
			_ = e.barrier.Delete(ctx, prefix+p)
		}
	}

	delete(e.users, username)

	return &engine.Response{Data: map[string]interface{}{
		"ok": true,
	}}, nil
}

// --- Internal helpers ---

// createUser generates a new keypair and stores it. Caller must hold e.mu write lock.
func (e *UserEngine) createUser(ctx context.Context, username string, autoProvisioned bool) (*userState, error) {
	priv, err := generateKey(e.config.KeyAlgorithm)
	if err != nil {
		return nil, fmt.Errorf("generate key for %s: %w", username, err)
	}

	if err := e.storeUserKeys(ctx, username, priv, autoProvisioned); err != nil {
		return nil, err
	}

	u := &userState{
		privKey: priv,
		pubKey:  priv.PublicKey(),
		config: &UserKeyConfig{
			Algorithm:       e.config.KeyAlgorithm,
			CreatedAt:       time.Now().UTC(),
			AutoProvisioned: autoProvisioned,
		},
	}
	e.users[username] = u
	return u, nil
}

// storeUserKeys persists user key material to the barrier. Caller must hold e.mu write lock.
func (e *UserEngine) storeUserKeys(ctx context.Context, username string, priv *ecdh.PrivateKey, autoProvisioned bool) error {
	prefix := e.mountPath + "users/" + username + "/"

	// Store private key.
	if err := e.barrier.Put(ctx, prefix+"priv.key", priv.Bytes()); err != nil {
		return fmt.Errorf("store private key: %w", err)
	}

	// Store public key.
	if err := e.barrier.Put(ctx, prefix+"pub.key", priv.PublicKey().Bytes()); err != nil {
		return fmt.Errorf("store public key: %w", err)
	}

	// Store config.
	cfg := &UserKeyConfig{
		Algorithm:       e.config.KeyAlgorithm,
		CreatedAt:       time.Now().UTC(),
		AutoProvisioned: autoProvisioned,
	}
	cfgData, err := json.Marshal(cfg)
	if err != nil {
		return fmt.Errorf("marshal key config: %w", err)
	}
	if err := e.barrier.Put(ctx, prefix+"config.json", cfgData); err != nil {
		return fmt.Errorf("store key config: %w", err)
	}

	return nil
}

// loadUser loads a single user's keys from the barrier into memory. Caller must hold e.mu write lock.
func (e *UserEngine) loadUser(ctx context.Context, username string) error {
	prefix := e.mountPath + "users/" + username + "/"

	// Load config first to know the algorithm.
	cfgData, err := e.barrier.Get(ctx, prefix+"config.json")
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}
	var cfg UserKeyConfig
	if err := json.Unmarshal(cfgData, &cfg); err != nil {
		return fmt.Errorf("parse config: %w", err)
	}

	// Load private key.
	privBytes, err := e.barrier.Get(ctx, prefix+"priv.key")
	if err != nil {
		return fmt.Errorf("load private key: %w", err)
	}

	curve, err := curveForAlgorithm(cfg.Algorithm)
	if err != nil {
		return err
	}

	priv, err := curve.NewPrivateKey(privBytes)
	if err != nil {
		return fmt.Errorf("parse private key: %w", err)
	}

	e.users[username] = &userState{
		privKey: priv,
		pubKey:  priv.PublicKey(),
		config:  &cfg,
	}
	return nil
}

// --- Cryptographic helpers ---

func generateKey(algorithm string) (*ecdh.PrivateKey, error) {
	curve, err := curveForAlgorithm(algorithm)
	if err != nil {
		return nil, err
	}
	return curve.GenerateKey(rand.Reader)
}

func curveForAlgorithm(algorithm string) (ecdh.Curve, error) {
	switch algorithm {
	case "x25519":
		return ecdh.X25519(), nil
	case "ecdh-p256":
		return ecdh.P256(), nil
	case "ecdh-p384":
		return ecdh.P384(), nil
	default:
		return nil, fmt.Errorf("%w: %s", ErrInvalidAlgorithm, algorithm)
	}
}

func validateKeyAlgorithm(alg string) error {
	_, err := curveForAlgorithm(alg)
	return err
}

func wrapDEKForRecipient(senderPriv *ecdh.PrivateKey, recipientPub *ecdh.PublicKey, dek []byte, sender, recipient string) (*recipientEntry, error) {
	shared, err := senderPriv.ECDH(recipientPub)
	if err != nil {
		return nil, fmt.Errorf("ECDH: %w", err)
	}
	defer crypto.Zeroize(shared)

	// Generate HKDF salt.
	salt := make([]byte, 32)
	if _, err := rand.Read(salt); err != nil {
		return nil, fmt.Errorf("generate salt: %w", err)
	}

	// Derive wrapping key.
	info := []byte(hkdfInfoPrefix + sender + ":" + recipient)
	wrappingKey := make([]byte, keySize)
	hkdfReader := hkdf.New(sha256.New, shared, salt, info)
	if _, err := io.ReadFull(hkdfReader, wrappingKey); err != nil {
		return nil, fmt.Errorf("HKDF: %w", err)
	}
	defer crypto.Zeroize(wrappingKey)

	// Wrap DEK with AES-256-GCM.
	wrapped, err := encryptAESGCM(wrappingKey, dek, nil)
	if err != nil {
		return nil, fmt.Errorf("wrap DEK: %w", err)
	}

	return &recipientEntry{
		Salt:       base64.StdEncoding.EncodeToString(salt),
		WrappedDEK: base64.StdEncoding.EncodeToString(wrapped),
	}, nil
}

func unwrapDEK(callerPriv *ecdh.PrivateKey, senderPub *ecdh.PublicKey, entry *recipientEntry, sender, caller string) ([]byte, error) {
	shared, err := callerPriv.ECDH(senderPub)
	if err != nil {
		return nil, fmt.Errorf("ECDH: %w", err)
	}
	defer crypto.Zeroize(shared)

	salt, err := base64.StdEncoding.DecodeString(entry.Salt)
	if err != nil {
		return nil, fmt.Errorf("decode salt: %w", err)
	}

	info := []byte(hkdfInfoPrefix + sender + ":" + caller)
	wrappingKey := make([]byte, keySize)
	hkdfReader := hkdf.New(sha256.New, shared, salt, info)
	if _, err := io.ReadFull(hkdfReader, wrappingKey); err != nil {
		return nil, fmt.Errorf("HKDF: %w", err)
	}
	defer crypto.Zeroize(wrappingKey)

	wrapped, err := base64.StdEncoding.DecodeString(entry.WrappedDEK)
	if err != nil {
		return nil, fmt.Errorf("decode wrapped DEK: %w", err)
	}

	dek, err := decryptAESGCM(wrappingKey, wrapped, nil)
	if err != nil {
		return nil, fmt.Errorf("unwrap DEK: %w", err)
	}
	return dek, nil
}

func encryptAESGCM(key, plaintext, aad []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, nonceSize)
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}
	ct := gcm.Seal(nil, nonce, plaintext, aad)
	// Return nonce + ciphertext+tag.
	result := make([]byte, nonceSize+len(ct))
	copy(result, nonce)
	copy(result[nonceSize:], ct)
	return result, nil
}

func decryptAESGCM(key, data, aad []byte) ([]byte, error) {
	if len(data) < nonceSize+16 { // nonce + at least one AES block
		return nil, ErrInvalidEnvelope
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonce := data[:nonceSize]
	ct := data[nonceSize:]
	return gcm.Open(nil, nonce, ct, aad)
}

func parseEnvelope(b64 string) (*envelope, error) {
	data, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return nil, fmt.Errorf("%w: base64 decode: %s", ErrInvalidEnvelope, err)
	}
	var env envelope
	if err := json.Unmarshal(data, &env); err != nil {
		return nil, fmt.Errorf("%w: json unmarshal: %s", ErrInvalidEnvelope, err)
	}
	if env.Version != 1 {
		return nil, fmt.Errorf("%w: unsupported version %d", ErrInvalidEnvelope, env.Version)
	}
	if env.Sender == "" || len(env.Recipients) == 0 {
		return nil, fmt.Errorf("%w: missing sender or recipients", ErrInvalidEnvelope)
	}
	return &env, nil
}

func extractRecipients(data map[string]interface{}) ([]string, error) {
	raw, ok := data["recipients"]
	if !ok {
		return nil, nil
	}

	switch v := raw.(type) {
	case []interface{}:
		names := make([]string, 0, len(v))
		for _, item := range v {
			s, ok := item.(string)
			if !ok {
				return nil, fmt.Errorf("user: recipient must be a string")
			}
			names = append(names, s)
		}
		return names, nil
	case []string:
		return v, nil
	default:
		return nil, fmt.Errorf("user: invalid recipients format")
	}
}