// Package session provides the post-handshake AEAD-protected Frame exchange and the sliding
// replay window — direct port of crates/aura-proto/src/session.rs.
package session

import "fmt"

// ReplayWindow is the width (in records) of the anti-replay sliding window.
const ReplayWindow uint64 = 64

// ErrReplay is returned when a record's sequence number is a duplicate or too old.
type ErrReplay struct{ Seq uint64 }

func (e *ErrReplay) Error() string { return fmt.Sprintf("aura/session: replay seq=%d", e.Seq) }

// Replay tracks the highest accepted sequence number and a 64-bit bitmap of the positions
// below it that have already been accepted. A datagram is accepted iff its seq is strictly
// newer than everything seen, or falls inside the window and was not previously seen.
type Replay struct {
	highest uint64
	bitmap  uint64
	seeded  bool
}

// NewReplay primes a window so the first expected record is `start` (the AEAD counter at the
// end of the handshake). Anything strictly below `start` is treated as already-consumed.
//
// This mirrors ReplayWindow::new in the Rust impl: highest = start - 1 (saturating),
// seeded = start > 0.
func NewReplay(start uint64) *Replay {
	r := &Replay{seeded: start > 0}
	if start > 0 {
		r.highest = start - 1
	}
	return r
}

// CheckAndSet records a seen seq. Returns nil if it is fresh; *ErrReplay otherwise.
func (r *Replay) CheckAndSet(seq uint64) error {
	if !r.seeded {
		// First-ever record (only reachable when started at 0): accept and seed.
		r.seeded = true
		r.highest = seq
		r.bitmap = 0
		return nil
	}
	if seq > r.highest {
		shift := seq - r.highest
		if shift >= 64 {
			r.bitmap = 0
		} else {
			r.bitmap = (r.bitmap << shift) | (1 << (shift - 1))
		}
		r.highest = seq
		return nil
	}
	// seq <= highest: must be inside the window and previously unseen.
	offset := r.highest - seq
	if offset >= ReplayWindow {
		return &ErrReplay{Seq: seq}
	}
	if offset == 0 {
		return &ErrReplay{Seq: seq}
	}
	mask := uint64(1) << (offset - 1)
	if r.bitmap&mask != 0 {
		return &ErrReplay{Seq: seq}
	}
	r.bitmap |= mask
	return nil
}