package session

import (
	"encoding/binary"
	"errors"
	"fmt"

	"github.com/aura/singbox-aura/aura/crypto"
	"github.com/aura/singbox-aura/aura/frame"
)

// SeqLen is the size of the per-record sequence-number prefix.
const SeqLen = 8

// PostHandshakeCounter is the AEAD counter at which the first application Data record starts,
// because each direction sealed exactly two encrypted handshake messages before it.
const PostHandshakeCounter uint64 = 2

// DatagramSender holds the outbound explicit-nonce AEAD plus the next sequence number to
// stamp. Produced by Session.IntoDatagramParts() after the handshake completes.
type DatagramSender struct {
	key *crypto.AeadKey
	seq uint64
}

// NewDatagramSender wraps a 32-byte key starting at the given counter.
func NewDatagramSender(rawKey []byte, startCounter uint64) (*DatagramSender, error) {
	k, err := crypto.NewAeadKey(rawKey)
	if err != nil {
		return nil, err
	}
	return &DatagramSender{key: k, seq: startCounter}, nil
}

// Seal encodes f, seals it under the next sequence number, and returns the on-wire datagram
// payload: seq(8 BE) || ciphertext.
func (s *DatagramSender) Seal(f *frame.Frame) []byte {
	seq := s.seq
	enc := frame.EncodeFrame(f)
	var seqBE [SeqLen]byte
	binary.BigEndian.PutUint64(seqBE[:], seq)
	ct := s.key.Seal(seq, enc, seqBE[:])
	out := make([]byte, 0, SeqLen+len(ct))
	out = append(out, seqBE[:]...)
	out = append(out, ct...)
	s.seq++
	return out
}

// NextSeq is the sequence number the next Seal will use (test/diagnostic helper).
func (s *DatagramSender) NextSeq() uint64 { return s.seq }

// DatagramReceiver authenticates, replay-checks, and decodes incoming datagram payloads.
type DatagramReceiver struct {
	key    *crypto.AeadKey
	replay *Replay
}

// NewDatagramReceiver wraps a 32-byte key plus a replay window primed at startCounter.
func NewDatagramReceiver(rawKey []byte, startCounter uint64) (*DatagramReceiver, error) {
	k, err := crypto.NewAeadKey(rawKey)
	if err != nil {
		return nil, err
	}
	return &DatagramReceiver{key: k, replay: NewReplay(startCounter)}, nil
}

// Open parses one datagram payload, runs the replay check first (so a duplicate cannot advance
// the AEAD state), then verifies and decodes the inner Frame.
func (r *DatagramReceiver) Open(datagram []byte) (*frame.Frame, error) {
	if len(datagram) < SeqLen {
		return nil, fmt.Errorf("aura/session: datagram shorter than seq prefix")
	}
	seqBE := datagram[:SeqLen]
	seq := binary.BigEndian.Uint64(seqBE)
	ct := datagram[SeqLen:]
	if err := r.replay.CheckAndSet(seq); err != nil {
		return nil, err
	}
	pt, err := r.key.Open(seq, ct, seqBE)
	if err != nil {
		return nil, err
	}
	return frame.DecodeFrame(pt)
}

// ErrUnexpectedMsg is returned by the stream half when the wire carries a non-Data record.
var ErrUnexpectedMsg = errors.New("aura/session: unexpected message type")