// Package crypto provides encryption-at-rest for messages, contacts, and
// calendar data. Every record gets its own key, derived from the master key
// via HKDF — compromising one encrypted file never exposes the master key or
// any other record.
package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"

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

const (
	keySize   = 32 // AES-256
	nonceSize = 12 // GCM standard nonce size
)

// MasterKey holds the decoded master key(s) used to derive per-record keys.
// Two keys allow decrypting old data during a key-rotation window: Current is
// used for new writes, Previous (if set) is tried as a fallback on decrypt.
type MasterKey struct {
	Current  []byte
	Previous []byte // nil if not rotating
}

// LoadMasterKey decodes hex-encoded master key(s) from config/env values.
func LoadMasterKey(currentHex, previousHex string) (*MasterKey, error) {
	cur, err := decodeKey(currentHex)
	if err != nil {
		return nil, fmt.Errorf("master key: %w", err)
	}
	mk := &MasterKey{Current: cur}
	if previousHex != "" {
		prev, err := decodeKey(previousHex)
		if err != nil {
			return nil, fmt.Errorf("previous master key: %w", err)
		}
		mk.Previous = prev
	}
	return mk, nil
}

func decodeKey(hexStr string) ([]byte, error) {
	b, err := hex.DecodeString(hexStr)
	if err != nil {
		return nil, fmt.Errorf("invalid hex: %w", err)
	}
	if len(b) != keySize {
		return nil, fmt.Errorf("must be %d bytes (%d hex chars), got %d bytes", keySize, keySize*2, len(b))
	}
	return b, nil
}

// deriveKey produces a per-record 32-byte key from the master key using HKDF-SHA256.
// recordID should be a stable, unique identifier for the record (e.g. message ID,
// contact UID) — using the same recordID always derives the same key, which is
// required for decryption to work.
func deriveKey(master []byte, recordID string, purpose string) ([]byte, error) {
	info := []byte(purpose + ":" + recordID)
	r := hkdf.New(sha256.New, master, nil, info)
	key := make([]byte, keySize)
	if _, err := io.ReadFull(r, key); err != nil {
		return nil, fmt.Errorf("hkdf derive: %w", err)
	}
	return key, nil
}

// Encrypt encrypts plaintext with a key derived from the master key and
// recordID. purpose namespaces the derivation (e.g. "message", "contact",
// "calendar", "dkim-key") so the same recordID used for different data types
// never collides. Output format: [12-byte nonce][ciphertext][16-byte GCM tag].
func Encrypt(mk *MasterKey, recordID, purpose string, plaintext []byte) ([]byte, error) {
	key, err := deriveKey(mk.Current, recordID, purpose)
	if err != nil {
		return nil, err
	}

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

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

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

// Decrypt decrypts data produced by Encrypt. It tries the current master key
// first, then falls back to the previous key (if set) — this lets records
// written before a key rotation still be read without a bulk re-encryption
// pass; callers should re-encrypt with the current key on next write if a
// fallback decrypt succeeds.
func Decrypt(mk *MasterKey, recordID, purpose string, data []byte) ([]byte, error) {
	if len(data) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}
	nonce, ciphertext := data[:nonceSize], data[nonceSize:]

	if pt, err := decryptWith(mk.Current, recordID, purpose, nonce, ciphertext); err == nil {
		return pt, nil
	}

	if mk.Previous != nil {
		if pt, err := decryptWith(mk.Previous, recordID, purpose, nonce, ciphertext); err == nil {
			return pt, nil
		}
	}

	return nil, fmt.Errorf("decrypt failed with current%s key",
		map[bool]string{true: " and previous", false: ""}[mk.Previous != nil])
}

func decryptWith(master []byte, recordID, purpose string, nonce, ciphertext []byte) ([]byte, error) {
	key, err := deriveKey(master, recordID, purpose)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	return gcm.Open(nil, nonce, ciphertext, nil)
}

// NeedsReencryption reports whether data was decrypted using the previous
// (not current) master key — callers use this to trigger lazy re-encryption
// on read during a key rotation window.
func NeedsReencryption(mk *MasterKey, recordID, purpose string, data []byte) bool {
	if mk.Previous == nil || len(data) < nonceSize {
		return false
	}
	nonce, ciphertext := data[:nonceSize], data[nonceSize:]
	if _, err := decryptWith(mk.Current, recordID, purpose, nonce, ciphertext); err == nil {
		return false // current key works fine
	}
	_, err := decryptWith(mk.Previous, recordID, purpose, nonce, ciphertext)
	return err == nil
}

// GenerateMasterKeyHex is a convenience helper for CLI tooling / setup docs —
// produces a fresh random master key as hex, ready to paste into GOMAIL_MASTER_KEY.
func GenerateMasterKeyHex() (string, error) {
	b := make([]byte, keySize)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}