package crypto

// Gnarl wire format for authenticated encrypted packets.
//
// Packet layout (64-byte header + variable ciphertext):
//
//	[10 bytes]  Nonce (counter-based, 80-bit)
//	[27 bytes]  Identity (GnarlMid of sender spore fingerprint)
//	[27 bytes]  Auth tag (GMid(per_packet_mac_key || nonce || identity || ciphertext_hash))
//	[variable]  Ciphertext (ChaCha20, keystream starts at block 1)
//
// Total overhead: 64 bytes. Fits within one UDP MTU fragment with ~1400
// bytes of payload.
//
// The auth tag binds the nonce, identity, and ciphertext together under
// the shared secret (derived from ExchangeV2). An attacker without the
// secret cannot forge the tag, reorder packets, or substitute identities.
//
// The identity field is the GnarlMid of the sender's spore fingerprint,
// giving a 27-byte (216-bit, birthday ~108-bit) collision-resistant
// identifier. This is NOT encrypted — it allows the receiver to look up
// the shared secret for this sender before decrypting.

import (
	"crypto/sha256"
	"encoding/binary"
	"errors"

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

const (
	// GnarlNonceLen is the nonce size in bytes (80 bits).
	GnarlNonceLen = 10

	// GnarlHeaderLen is the total header size: nonce + identity + auth tag.
	GnarlHeaderLen = GnarlNonceLen + GnarlMidBytes + GnarlMidBytes // 10 + 27 + 27 = 64
)

// GnarlPacket is an authenticated encrypted packet in the Gnarl wire format.
type GnarlPacket struct {
	Nonce      [GnarlNonceLen]byte // 80-bit nonce (counter-based)
	Identity   GnarlMid            // sender identity (GMid of spore fingerprint)
	AuthTag    GnarlMid            // authentication tag
	Ciphertext []byte              // CTR-encrypted payload
}

// GnarlSeal encrypts plaintext and produces an authenticated packet.
//
// Parameters:
//   - secret: shared Hamadryad secret from key exchange
//   - identity: sender's GnarlMid (public, used for session lookup)
//   - nonce: 80-bit nonce (must be unique per packet for a given secret)
//   - plaintext: data to encrypt
//
// The ciphertext is produced by ChaCha20 keyed with the shared secret.
// The auth tag uses a per-packet MAC key derived from ChaCha20 block 0,
// eliminating SWIFFT homomorphic relations between tags across packets.
func GnarlSeal(secret Hamadryad, identity GnarlMid, nonce [GnarlNonceLen]byte, plaintext []byte) *GnarlPacket {
	macKey, cipher := gnarlCipherAndMAC(secret, nonce)

	// Encrypt: ChaCha20 keystream (starts at block 1, after MAC key block).
	ciphertext := make([]byte, len(plaintext))
	cipher.XORKeyStream(ciphertext, plaintext)

	// Auth tag with per-packet MAC key.
	authTag := gnarlAuthTag(macKey, nonce, identity, ciphertext)

	return &GnarlPacket{
		Nonce:      nonce,
		Identity:   identity,
		AuthTag:    authTag,
		Ciphertext: ciphertext,
	}
}

// GnarlOpen verifies and decrypts an authenticated packet.
//
// Returns the plaintext if the auth tag is valid, or an error if
// verification fails (tampered, wrong key, replayed, etc.).
func GnarlOpen(secret Hamadryad, pkt *GnarlPacket) ([]byte, error) {
	if pkt == nil {
		return nil, errors.New("gnarl: nil packet")
	}

	macKey, cipher := gnarlCipherAndMAC(secret, pkt.Nonce)

	// Verify auth tag with per-packet MAC key.
	expected := gnarlAuthTag(macKey, pkt.Nonce, pkt.Identity, pkt.Ciphertext)
	if expected != pkt.AuthTag {
		return nil, errors.New("gnarl: authentication failed")
	}

	// Decrypt (cipher is already positioned at block 1).
	plaintext := make([]byte, len(pkt.Ciphertext))
	cipher.XORKeyStream(plaintext, pkt.Ciphertext)
	return plaintext, nil
}

// MarshalGnarlPacket serializes a packet to wire format.
func MarshalGnarlPacket(pkt *GnarlPacket) []byte {
	buf := make([]byte, GnarlHeaderLen+len(pkt.Ciphertext))
	copy(buf[0:GnarlNonceLen], pkt.Nonce[:])
	copy(buf[GnarlNonceLen:GnarlNonceLen+GnarlMidBytes], pkt.Identity[:])
	copy(buf[GnarlNonceLen+GnarlMidBytes:GnarlHeaderLen], pkt.AuthTag[:])
	copy(buf[GnarlHeaderLen:], pkt.Ciphertext)
	return buf
}

// UnmarshalGnarlPacket deserializes a packet from wire format.
func UnmarshalGnarlPacket(data []byte) (*GnarlPacket, error) {
	if len(data) < GnarlHeaderLen {
		return nil, errors.New("gnarl: packet too short")
	}

	pkt := &GnarlPacket{}
	copy(pkt.Nonce[:], data[0:GnarlNonceLen])
	copy(pkt.Identity[:], data[GnarlNonceLen:GnarlNonceLen+GnarlMidBytes])
	copy(pkt.AuthTag[:], data[GnarlNonceLen+GnarlMidBytes:GnarlHeaderLen])
	pkt.Ciphertext = make([]byte, len(data)-GnarlHeaderLen)
	copy(pkt.Ciphertext, data[GnarlHeaderLen:])
	return pkt, nil
}

// gnarlCipherAndMAC derives a per-packet MAC key and a ChaCha20 cipher
// positioned for encryption, following the RFC 8439 pattern:
//   - Block 0: consumed to derive a one-time MAC key (56 bytes)
//   - Block 1+: available for encryption keystream
//
// The ChaCha20 key is derived from the shared secret via a single
// Hamadryad hash evaluation (one-time key derivation, not a PRF).
func gnarlCipherAndMAC(secret Hamadryad, nonce [GnarlNonceLen]byte) (Hamadryad, *chacha20.Cipher) {
	// Derive 32-byte ChaCha20 key from shared secret.
	keyHash := Hash(append([]byte("gnarl-chacha20-key"), secret[:]...))
	var key [32]byte
	copy(key[:], keyHash[:32])

	// Zero-pad 10-byte wire nonce to 12 bytes.
	var chachaNonce [chacha20.NonceSize]byte
	copy(chachaNonce[:GnarlNonceLen], nonce[:])

	cipher, _ := chacha20.NewUnauthenticatedCipher(key[:], chachaNonce[:])

	// Derive per-packet MAC key from block 0 (first 64 bytes of keystream).
	// Use first 56 bytes (Hamadryad-sized) as the one-time MAC key.
	var macBlock [64]byte
	cipher.XORKeyStream(macBlock[:], macBlock[:])
	var macKey Hamadryad
	copy(macKey[:], macBlock[:HamBytes])

	// Cipher is now positioned at block 1, ready for encryption.
	return macKey, cipher
}

// gnarlAuthTag computes the authentication tag for a packet.
// tag = GMid(macKey || nonce || identity || GHash(ciphertext))
//
// The macKey is a per-packet one-time key derived from ChaCha20 block 0,
// so each packet's tag uses an independent pseudorandom key. This
// eliminates homomorphic relations between tags across packets.
func gnarlAuthTag(macKey Hamadryad, nonce [GnarlNonceLen]byte, identity GnarlMid, ciphertext []byte) GnarlMid {
	// Hash the ciphertext with GHash for a fixed-size binding.
	ctHash := GHash(ciphertext)

	var tagInput []byte
	tagInput = append(tagInput, macKey[:]...)
	tagInput = append(tagInput, nonce[:]...)
	tagInput = append(tagInput, identity[:]...)
	tagInput = append(tagInput, ctHash[:]...)
	return GMid(tagInput)
}

// GnarlSchnorrChallenge computes the Fiat-Shamir challenge for Gnarl
// Schnorr signatures using SHA-256 truncated to 27 bytes.
//
// SHA-256 is used instead of GMid because the Schnorr security proof
// requires random oracle behaviour. SWIFFT's additive homomorphism
// makes it distinguishable from a random oracle, which would be
// exploitable in multi-signature protocols. SHA-256 has no known
// algebraic structure and is the standard choice for Fiat-Shamir.
//
// The output is truncated to 27 bytes (216 bits) to match the Gnarl
// scalar field size. Collision resistance: 2^108 (birthday bound on
// 216-bit output), which exceeds the ~107-bit Pollard-rho security
// of the Gnarl torus discrete log.
func GnarlSchnorrChallenge(data []byte) [GnarlMidBytes]byte {
	h := sha256.Sum256(data)
	var result [GnarlMidBytes]byte
	copy(result[:], h[:GnarlMidBytes])
	return result
}

// GnarlNonceFromCounter creates a nonce from a 64-bit counter and 16-bit epoch.
// This covers the common case of counter-based nonces with epoch rotation.
func GnarlNonceFromCounter(epoch uint16, counter uint64) [GnarlNonceLen]byte {
	var nonce [GnarlNonceLen]byte
	binary.LittleEndian.PutUint64(nonce[0:8], counter)
	binary.LittleEndian.PutUint16(nonce[8:10], epoch)
	return nonce
}