// Key wrapping for cold-start unlock.
//
// The at-rest AES-256 key is wrapped with a key derived from the passkey
// credential public key (stable across assertions) via HKDF-SHA256, then
// AES-256-GCM. The wrapped blob is stored on disk; at cold-start the passkey
// assertion provides the credential public key to unwrap it.
//
// The passkey credential is a P-256 ECDSA public key. Its raw uncompressed
// bytes (65 bytes, 0x04 || X || Y) are the HKDF input — high-entropy, stable.
//
// Blob format: salt (16) || nonce (12) || AES-256-GCM ciphertext.
// No file magic — the caller (mavend) owns the file path.
package webauthn

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

const (
	// saltLen — HKDF salt length. 16 bytes is standard.
	saltLen = 16
	// nonceLen — AES-GCM standard nonce length.
	nonceLen = 12
	// keyLen — AES-256 key length.
	keyLen = 32
	// wrapInfo — HKDF info string for domain separation.
	wrapInfo = "maven-passkey-keywrap-v1"
)

var (
	ErrKeyWrap     = errors.New("webauthn: key wrap failed")
	ErrKeyUnwrap   = errors.New("webauthn: key unwrap failed (wrong credential?)")
	ErrBlobTooLong = errors.New("webauthn: wrapped blob too long")
)

// WrapKey derives a wrapping key from credPublicKey via HKDF-SHA256 and
// AES-GCM-wraps plaintextKey. Returns the blob: salt || nonce || ciphertext.
// plaintextKey must be exactly 32 bytes (AES-256).
func WrapKey(plaintextKey, credPublicKey []byte) ([]byte, error) {
	if len(plaintextKey) != keyLen {
		return nil, fmt.Errorf("%w: plaintext key must be %d bytes", ErrKeyWrap, keyLen)
	}
	if len(credPublicKey) == 0 {
		return nil, fmt.Errorf("%w: empty credential public key", ErrKeyWrap)
	}

	salt := make([]byte, saltLen)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, fmt.Errorf("%w: salt: %v", ErrKeyWrap, err)
	}

	wrapKey := hkdfSHA256(credPublicKey, salt, []byte(wrapInfo), keyLen)

	nonce := make([]byte, nonceLen)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("%w: nonce: %v", ErrKeyWrap, err)
	}

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

	// Seal appends ciphertext+tag to nonce (which becomes nonce||ct).
	ct := gcm.Seal(nil, nonce, plaintextKey, nil)

	out := make([]byte, 0, saltLen+nonceLen+len(ct))
	out = append(out, salt...)
	out = append(out, nonce...)
	out = append(out, ct...)
	return out, nil
}

// UnwrapKey extracts the salt from blob, re-derives the wrapping key from
// credPublicKey, and AES-GCM-unwraps. Returns the plaintext 32-byte AES key.
func UnwrapKey(blob, credPublicKey []byte) ([]byte, error) {
	if len(blob) < saltLen+nonceLen+1 {
		return nil, fmt.Errorf("%w: blob too short (%d)", ErrKeyUnwrap, len(blob))
	}
	if len(blob) > 1<<20 { // 1MB sanity limit
		return nil, ErrBlobTooLong
	}
	if len(credPublicKey) == 0 {
		return nil, fmt.Errorf("%w: empty credential public key", ErrKeyUnwrap)
	}

	salt := blob[:saltLen]
	nonce := blob[saltLen : saltLen+nonceLen]
	ct := blob[saltLen+nonceLen:]

	wrapKey := hkdfSHA256(credPublicKey, salt, []byte(wrapInfo), keyLen)

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

	plain, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		return nil, fmt.Errorf("%w: decrypt failed (wrong credential?)", ErrKeyUnwrap)
	}
	return plain, nil
}

// hkdfSHA256 implements HKDF-SHA256 (RFC 5869) using only stdlib.
//
// Input:
//   - secret: the input key material (credential public key bytes)
//   - salt: random salt (16 bytes)
//   - info: optional context string for domain separation
//   - length: desired output length in bytes
//
// Output: length bytes of derived key material.
//
// HKDF is extract-then-expand. We use HMAC-SHA256 for both steps. This avoids
// importing golang.org/x/crypto/hkdf — a ~30-line function vs a new dep. The
// tradeoff is no constant-time guarantees on the extract step beyond HMAC's;
// acceptable here because the input is already high-entropy key material (a
// P-256 public key), not a low-entropy passphrase.
func hkdfSHA256(secret, salt, info []byte, length int) []byte {
	// Step 1: Extract — PRK = HMAC-SHA256(salt, secret)
	// If salt is nil/empty, use a zero-filled block (RFC 5869 §2.2).
	if salt == nil {
		salt = make([]byte, sha256.Size)
	}
	mac := hmac.New(sha256.New, salt)
	mac.Write(secret)
	prk := mac.Sum(nil)

	// Step 2: Expand — produce length bytes via T(i) = HMAC-SHA256(PRK, T(i-1) || info || i)
	// Where T(0) = empty, i is a byte counter starting at 1.
	out := make([]byte, 0, length)
	block := make([]byte, 0, sha256.Size+len(info)+1)
	var t []byte // T(i-1)
	for counter := byte(1); len(out) < length; counter++ {
		block = block[:0]
		block = append(block, t...)
		block = append(block, info...)
		block = append(block, counter)

		mac.Reset()
		mac.Write(block)
		t = mac.Sum(prk[:0]) // reuse prk buffer — mac.Sum appends to its arg
		// t now starts with prk[:0] (empty) followed by the HMAC result.
		// Since we need just the HMAC result (sha256.Size bytes), re-slice.
		t = t[len(t)-sha256.Size:]
		out = append(out, t...)
	}
	return out[:length]
}

// encodeUint32 — big-endian uint32 for the blob format header, if needed.
func encodeUint32(v uint32) []byte {
	var b [4]byte
	binary.BigEndian.PutUint32(b[:], v)
	return b[:]
}