package crypto

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

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

const (
	InfoString      = "aes-demo-v1"
	AlgorithmName   = "AES-256-GCM-HKDF-SHA256"
	KeyK1Length     = 32
	KdfSaltLength   = 16
	GcmNonceLength  = 12
)

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
}

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
}