package crypto

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

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

const (
	// InfoString is used to bind HKDF output to this application/version.
	InfoString = "aes-demo-v1"
	// AlgorithmName is sent alongside payloads to assert parameters.
	AlgorithmName = "AES-256-GCM-HKDF-SHA256"

	// KeyK1Length is the required length of K1 in bytes.
	KeyK1Length = 32
	// KdfSaltLength is the random nonce length used as HKDF salt.
	KdfSaltLength = 16
	// GcmNonceLength is the AEAD GCM nonce length (12 bytes recommended).
	GcmNonceLength = 12
)

// DeriveKeyAndGcmNonce derives the encryption key K2 and the GCM nonce using HKDF-SHA256.
// secret = K1, salt = kdfSalt (aka transport nonce), info = InfoString.
func DeriveKeyAndGcmNonce(k1 []byte, kdfSalt []byte) (k2 []byte, gcmNonce []byte, err error) {
	if len(k1) != KeyK1Length {
		return nil, nil, fmt.Errorf("k1 must be %d bytes, got %d", KeyK1Length, len(k1))
	}
	if len(kdfSalt) != KdfSaltLength {
		return nil, nil, fmt.Errorf("kdf salt (nonce) must be %d bytes, got %d", KdfSaltLength, len(kdfSalt))
	}

	h := hkdf.New(sha256.New, k1, kdfSalt, []byte(InfoString))
	buf := make([]byte, KeyK1Length+GcmNonceLength)
	if _, err := io.ReadFull(h, buf); err != nil {
		return nil, nil, fmt.Errorf("hkdf expand: %w", err)
	}
	return buf[:KeyK1Length], buf[KeyK1Length:], nil
}

// GenerateKdfSalt returns a fresh random 16-byte salt for HKDF.
func GenerateKdfSalt() ([]byte, error) {
	salt := make([]byte, KdfSaltLength)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, fmt.Errorf("rand salt: %w", err)
	}
	return salt, nil
}

// EncryptGCM encrypts plaintext using K1 and provided kdfSalt by deriving K2 and gcmNonce.
// Returns ciphertext (includes GCM tag).
func EncryptGCM(k1, kdfSalt, plaintext, aad []byte) ([]byte, error) {
	k2, gcmNonce, err := DeriveKeyAndGcmNonce(k1, kdfSalt)
	if err != nil {
		return nil, err
	}
	blk, err := aes.NewCipher(k2)
	if err != nil {
		return nil, fmt.Errorf("new cipher: %w", err)
	}
	a, err := cipher.NewGCMWithNonceSize(blk, GcmNonceLength)
	if err != nil {
		return nil, fmt.Errorf("new gcm: %w", err)
	}
	return a.Seal(nil, gcmNonce, plaintext, aad), nil
}

// DecryptGCM decrypts ciphertext using K1 and provided kdfSalt by deriving K2 and gcmNonce.
func DecryptGCM(k1, kdfSalt, ciphertext, aad []byte) ([]byte, error) {
	k2, gcmNonce, err := DeriveKeyAndGcmNonce(k1, kdfSalt)
	if err != nil {
		return nil, err
	}
	blk, err := aes.NewCipher(k2)
	if err != nil {
		return nil, fmt.Errorf("new cipher: %w", err)
	}
	a, err := cipher.NewGCMWithNonceSize(blk, GcmNonceLength)
	if err != nil {
		return nil, fmt.Errorf("new gcm: %w", err)
	}
	pt, err := a.Open(nil, gcmNonce, ciphertext, aad)
	if err != nil {
		return nil, errors.New("gcm open failed: authentication failed or corrupted data")
	}
	return pt, nil
}