// Package crypto provides Argon2id KDF, AES-256-GCM encryption, and key helpers.
package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/subtle"
	"errors"
	"fmt"

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

const (
	// KeySize is the size of AES-256 keys in bytes.
	KeySize = 32
	// NonceSize is the size of AES-GCM nonces in bytes.
	NonceSize = 12
	// SaltSize is the size of Argon2id salts in bytes.
	SaltSize = 32

	// BarrierVersion is the version byte prefix for encrypted barrier entries.
	BarrierVersion byte = 0x01

	// Default Argon2id parameters.
	DefaultArgon2Time    = 3
	DefaultArgon2Memory  = 128 * 1024 // 128 MiB in KiB
	DefaultArgon2Threads = 4
)

var (
	ErrInvalidCiphertext = errors.New("crypto: invalid ciphertext")
	ErrDecryptionFailed  = errors.New("crypto: decryption failed")
)

// Argon2Params holds Argon2id KDF parameters.
type Argon2Params struct {
	Time    uint32
	Memory  uint32 // in KiB
	Threads uint8
}

// DefaultArgon2Params returns the default Argon2id parameters.
func DefaultArgon2Params() Argon2Params {
	return Argon2Params{
		Time:    DefaultArgon2Time,
		Memory:  DefaultArgon2Memory,
		Threads: DefaultArgon2Threads,
	}
}

// DeriveKey derives a 256-bit key from password and salt using Argon2id.
func DeriveKey(password []byte, salt []byte, params Argon2Params) []byte {
	return argon2.IDKey(password, salt, params.Time, params.Memory, params.Threads, KeySize)
}

// GenerateKey generates a random 256-bit key.
func GenerateKey() ([]byte, error) {
	key := make([]byte, KeySize)
	if _, err := rand.Read(key); err != nil {
		return nil, fmt.Errorf("crypto: generate key: %w", err)
	}
	return key, nil
}

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

// Encrypt encrypts plaintext with AES-256-GCM using the given key.
// Returns: [version byte][12-byte nonce][ciphertext+tag]
func Encrypt(key, plaintext []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("crypto: new cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("crypto: new gcm: %w", err)
	}
	nonce := make([]byte, NonceSize)
	if _, err := rand.Read(nonce); err != nil {
		return nil, fmt.Errorf("crypto: generate nonce: %w", err)
	}
	ciphertext := gcm.Seal(nil, nonce, plaintext, nil)

	// Format: [version][nonce][ciphertext+tag]
	result := make([]byte, 1+NonceSize+len(ciphertext))
	result[0] = BarrierVersion
	copy(result[1:1+NonceSize], nonce)
	copy(result[1+NonceSize:], ciphertext)
	return result, nil
}

// Decrypt decrypts ciphertext produced by Encrypt.
func Decrypt(key, data []byte) ([]byte, error) {
	if len(data) < 1+NonceSize+aes.BlockSize {
		return nil, ErrInvalidCiphertext
	}
	if data[0] != BarrierVersion {
		return nil, fmt.Errorf("crypto: unsupported version: %d", data[0])
	}
	nonce := data[1 : 1+NonceSize]
	ciphertext := data[1+NonceSize:]

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("crypto: new cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("crypto: new gcm: %w", err)
	}
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, ErrDecryptionFailed
	}
	return plaintext, nil
}

// Zeroize overwrites a byte slice with zeros.
func Zeroize(b []byte) {
	for i := range b {
		b[i] = 0
	}
}

// ConstantTimeEqual compares two byte slices in constant time.
func ConstantTimeEqual(a, b []byte) bool {
	return subtle.ConstantTimeCompare(a, b) == 1
}