package crypto

import (
	"crypto/cipher"
	"encoding/binary"
	"fmt"

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

// NonceLen is the AEAD nonce length (96 bits for ChaCha20-Poly1305).
const NonceLen = 12

// NonceFor reproduces the AeadSession::nonce_for layout exactly:
//
//	nonce[0..8]  = LE(u64) counter
//	nonce[8..12] = 0
//
// Both stream- and datagram-mode AEADs share this nonce derivation; the only difference is
// whether the counter is advanced lock-step (stream) or carried on the wire (datagram).
func NonceFor(counter uint64) [NonceLen]byte {
	var n [NonceLen]byte
	binary.LittleEndian.PutUint64(n[0:8], counter)
	return n
}

// AeadKey wraps a 32-byte ChaCha20-Poly1305 key for explicit-nonce datagram use. The caller owns
// nonce uniqueness — Aura's datagram codec carries the counter on the wire as `seq`.
type AeadKey struct {
	aead cipher.AEAD
}

// NewAeadKey builds an AeadKey from a 32-byte key. Returns an error if the key is the wrong
// size; ChaCha20-Poly1305 always wants 32.
func NewAeadKey(key []byte) (*AeadKey, error) {
	if len(key) != SessionKeyLen {
		return nil, fmt.Errorf("aead key must be %d bytes, got %d", SessionKeyLen, len(key))
	}
	a, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, fmt.Errorf("chacha20poly1305.New: %w", err)
	}
	return &AeadKey{aead: a}, nil
}

// Seal encrypts plaintext under the nonce derived from counter, returning ciphertext||tag.
func (k *AeadKey) Seal(counter uint64, plaintext, aad []byte) []byte {
	nonce := NonceFor(counter)
	return k.aead.Seal(nil, nonce[:], plaintext, aad)
}

// Open authenticates and decrypts ciphertext (which must include the 16-byte Poly1305 tag).
// Returns the plaintext, or an error on authentication failure.
func (k *AeadKey) Open(counter uint64, ciphertext, aad []byte) ([]byte, error) {
	nonce := NonceFor(counter)
	out, err := k.aead.Open(nil, nonce[:], ciphertext, aad)
	if err != nil {
		return nil, fmt.Errorf("aead open: %w", err)
	}
	return out, nil
}

// AeadSession is the stream-mode counterpart: it holds the key plus a monotonically increasing
// 64-bit counter that advances on every Seal and Open. Used by the handshake's encrypted
// messages (ServerAuth, ClientAuth, Finished) so the two sides stay in lockstep without putting
// the counter on the wire.
type AeadSession struct {
	key     *AeadKey
	counter uint64
}

// NewAeadSession starts a session at counter 0.
func NewAeadSession(rawKey []byte) (*AeadSession, error) {
	k, err := NewAeadKey(rawKey)
	if err != nil {
		return nil, err
	}
	return &AeadSession{key: k, counter: 0}, nil
}

// Counter is the current counter (the nonce that the next Seal/Open will use). Test-only and
// used by Session.IntoDatagramParts to hand off the explicit-nonce key.
func (s *AeadSession) Counter() uint64 { return s.counter }

// Seal seals plaintext at the current counter then advances it.
func (s *AeadSession) Seal(plaintext, aad []byte) []byte {
	ct := s.key.Seal(s.counter, plaintext, aad)
	s.counter++
	return ct
}

// Open verifies+decrypts ciphertext at the current counter then advances it (symmetric to Seal
// so a failed decrypt keeps the two ends aligned).
func (s *AeadSession) Open(ciphertext, aad []byte) ([]byte, error) {
	pt, err := s.key.Open(s.counter, ciphertext, aad)
	s.counter++
	return pt, err
}

// IntoKey returns the underlying AeadKey so datagram-mode codecs can continue at the same
// counter without re-deriving anything (matches Rust's into_parts).
func (s *AeadSession) IntoKey() *AeadKey { return s.key }