neural-hunt/internal/runtime/state.go

package runtime

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"neuralhunt/internal/data"
)

var (
	ErrPresenceConflict = errors.New("identity already connected")
	ErrRateLimited      = errors.New("rate limited")
	ErrBadSequence      = errors.New("bad sequence")
)

type guessKey struct{ TaskID, ClientID string }

type presenceEntry struct {
	SessionID string
	LeaseID   uint64
}

type GuessState struct {
	NextSeq    int64
	LastGuess  time.Time
	BestScore  float64
	GuessCount int64
	PublicSeed string
	Revision   int64
}

type State struct {
	mu       sync.RWMutex
	presence map[string]presenceEntry
	selected map[string]string
	guesses  map[guessKey]GuessState

	started           time.Time
	presenceSeq       atomic.Uint64
	guessesTotal      atomic.Uint64
	improvementsTotal atomic.Uint64
	sqliteWritesTotal atomic.Uint64
	rejectedTotal     atomic.Uint64

	rateMu sync.Mutex
	sec    [60]rateBucket
}

type rateBucket struct {
	unix         int64
	guesses      uint64
	improvements uint64
	sqliteWrites uint64
	rejected     uint64
}

func New() *State {
	return &State{presence: make(map[string]presenceEntry), selected: make(map[string]string), guesses: make(map[guessKey]GuessState), started: time.Now()}
}

// AcquirePresence reserves a client identity for one authenticated browser/session.
// Reconnects from the SAME session are allowed and receive a new lease generation.
// The generation is critical: an old websocket may finish closing after a replacement
// connection has already been established. Its deferred ReleasePresence must not delete
// the newer connection's presence entry.
func (s *State) AcquirePresence(clientID, sessionID string) (uint64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if cur, ok := s.presence[clientID]; ok && cur.SessionID != sessionID {
		return 0, ErrPresenceConflict
	}
	lease := s.presenceSeq.Add(1)
	s.presence[clientID] = presenceEntry{SessionID: sessionID, LeaseID: lease}
	return lease, nil
}
func (s *State) HasPresence(clientID, sessionID string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	cur, ok := s.presence[clientID]
	return ok && cur.SessionID == sessionID
}
func (s *State) ReleasePresence(clientID, sessionID string, leaseID uint64) {
	s.mu.Lock()
	if cur, ok := s.presence[clientID]; ok && cur.SessionID == sessionID && cur.LeaseID == leaseID {
		delete(s.presence, clientID)
	}
	s.mu.Unlock()
}
func (s *State) IsConnected(clientID string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.presence[clientID]
	return ok
}
func (s *State) ConnectedCount() int64 {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return int64(len(s.presence))
}

// ForgetClient removes non-durable runtime state for an identity that was
// deleted by the admin profile cleanup tool. The caller must only pass clients
// that are currently disconnected.
func (s *State) ForgetClient(clientID string) {
	s.mu.Lock()
	delete(s.presence, clientID)
	delete(s.selected, clientID)
	for k := range s.guesses {
		if k.ClientID == clientID {
			delete(s.guesses, k)
		}
	}
	s.mu.Unlock()
}

func (s *State) SetTaskSelection(clientID, taskID string) {
	s.mu.Lock()
	if taskID == "" {
		delete(s.selected, clientID)
	} else {
		s.selected[clientID] = taskID
	}
	s.mu.Unlock()
}

func (s *State) TaskSelection(clientID string) string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.selected[clientID]
}

func (s *State) ReplaceTaskSelection(oldTaskID, newTaskID string) {
	s.mu.Lock()
	for cid, taskID := range s.selected {
		if taskID == oldTaskID {
			s.selected[cid] = newTaskID
		}
	}
	s.mu.Unlock()
}

func (s *State) InitGuess(task data.Task, cid string, persisted GuessState) GuessState {
	k := guessKey{task.ID, cid}
	s.mu.Lock()
	defer s.mu.Unlock()
	if cur, ok := s.guesses[k]; ok {
		// Preserve the hot-path sequence across admin changes that keep the same seed.
		// A reroll changes public_seed, which intentionally reloads the persisted reset.
		if cur.PublicSeed == task.PublicSeed {
			cur.Revision = task.Revision
			s.guesses[k] = cur
			return cur
		}
	}
	persisted.PublicSeed = task.PublicSeed
	persisted.Revision = task.Revision
	s.guesses[k] = persisted
	return persisted
}

func (s *State) Current(task data.Task, cid string) (GuessState, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	g, ok := s.guesses[guessKey{task.ID, cid}]
	if !ok || g.PublicSeed != task.PublicSeed {
		return GuessState{}, false
	}
	return g, true
}

type AcceptResult struct {
	State    GuessState
	Previous GuessState
	Improved bool
}

// CanSubmit validates sequence and per-client timing before a request is
// admitted into the global-per-task lottery. It does not mutate state, so an
// unselected ticket can be consumed explicitly by SkipLottery.
func (s *State) CanSubmit(task data.Task, cid string, seq int64, minInterval time.Duration) error {
	k := guessKey{task.ID, cid}
	now := time.Now().UTC()
	s.mu.RLock()
	g, ok := s.guesses[k]
	s.mu.RUnlock()
	if !ok || g.PublicSeed != task.PublicSeed || seq != g.NextSeq {
		s.record(false, false, true)
		return ErrBadSequence
	}
	if !g.LastGuess.IsZero() && minInterval > 0 && now.Sub(g.LastGuess) < minInterval {
		s.record(false, false, true)
		return ErrRateLimited
	}
	return nil
}

// SkipLottery consumes a valid sequence that lost the random draw without
// evaluating/scoring it. This prevents the same deterministic guess from being
// resubmitted forever while keeping guess_count reserved for actually accepted
// and evaluated tips.
func (s *State) SkipLottery(task data.Task, cid string, seq int64, minInterval time.Duration) (int64, error) {
	k := guessKey{task.ID, cid}
	now := time.Now().UTC()
	s.mu.Lock()
	g, ok := s.guesses[k]
	if !ok || g.PublicSeed != task.PublicSeed || seq != g.NextSeq {
		next := g.NextSeq
		s.mu.Unlock()
		s.record(false, false, true)
		return next, ErrBadSequence
	}
	if !g.LastGuess.IsZero() && minInterval > 0 && now.Sub(g.LastGuess) < minInterval {
		next := g.NextSeq
		s.mu.Unlock()
		s.record(false, false, true)
		return next, ErrRateLimited
	}
	g.NextSeq++
	g.LastGuess = now
	g.Revision = task.Revision
	s.guesses[k] = g
	next := g.NextSeq
	s.mu.Unlock()
	s.record(false, false, true)
	return next, nil
}

func (s *State) Accept(task data.Task, cid string, seq int64, score float64, minInterval time.Duration) (AcceptResult, error) {
	k := guessKey{task.ID, cid}
	now := time.Now().UTC()
	s.mu.Lock()
	g, ok := s.guesses[k]
	if !ok || g.PublicSeed != task.PublicSeed {
		s.mu.Unlock()
		return AcceptResult{}, ErrBadSequence
	}
	if seq != g.NextSeq {
		s.mu.Unlock()
		s.record(false, false, true)
		return AcceptResult{}, ErrBadSequence
	}
	prev := g
	if !g.LastGuess.IsZero() && minInterval > 0 && now.Sub(g.LastGuess) < minInterval {
		s.mu.Unlock()
		s.record(false, false, true)
		return AcceptResult{}, ErrRateLimited
	}
	g.NextSeq++
	g.GuessCount++
	g.LastGuess = now
	improved := score > g.BestScore
	if improved {
		g.BestScore = score
	}
	g.Revision = task.Revision
	s.guesses[k] = g
	s.mu.Unlock()
	s.record(true, improved, false)
	return AcceptResult{State: g, Previous: prev, Improved: improved}, nil
}

func (s *State) Restore(task data.Task, cid string, acceptedNextSeq int64, previous GuessState) {
	k := guessKey{task.ID, cid}
	s.mu.Lock()
	if cur, ok := s.guesses[k]; ok && cur.NextSeq == acceptedNextSeq && cur.PublicSeed == task.PublicSeed {
		s.guesses[k] = previous
	}
	s.mu.Unlock()
}

func (s *State) MarkSQLiteWrite() {
	s.sqliteWritesTotal.Add(1)
	s.rateMu.Lock()
	b := s.bucketLocked(time.Now().Unix())
	b.sqliteWrites++
	s.rateMu.Unlock()
}
func (s *State) record(guess, improvement, rejected bool) {
	if guess {
		s.guessesTotal.Add(1)
	}
	if improvement {
		s.improvementsTotal.Add(1)
	}
	if rejected {
		s.rejectedTotal.Add(1)
	}
	s.rateMu.Lock()
	b := s.bucketLocked(time.Now().Unix())
	if guess {
		b.guesses++
	}
	if improvement {
		b.improvements++
	}
	if rejected {
		b.rejected++
	}
	s.rateMu.Unlock()
}
func (s *State) bucketLocked(sec int64) *rateBucket {
	i := sec % 60
	if s.sec[i].unix != sec {
		s.sec[i] = rateBucket{unix: sec}
	}
	return &s.sec[i]
}

func (s *State) Rates(window int64) (guesses, improvements, sqliteWrites, rejected float64) {
	if window < 1 {
		window = 5
	}
	if window > 60 {
		window = 60
	}
	now := time.Now().Unix()
	var g, i, w, r uint64
	s.rateMu.Lock()
	for x := range s.sec {
		b := s.sec[x]
		if b.unix > now-window {
			g += b.guesses
			i += b.improvements
			w += b.sqliteWrites
			r += b.rejected
		}
	}
	s.rateMu.Unlock()
	d := float64(window)
	return float64(g) / d, float64(i) / d, float64(w) / d, float64(r) / d
}

type Metrics struct {
	Connected          int64   `json:"connected"`
	GuessesPerSec      float64 `json:"guesses_per_sec"`
	ImprovementsPerSec float64 `json:"improvements_per_sec"`
	SQLiteWritesPerSec float64 `json:"sqlite_writes_per_sec"`
	RejectedPerSec     float64 `json:"rejected_per_sec"`
	GuessesTotal       uint64  `json:"guesses_total"`
	ImprovementsTotal  uint64  `json:"improvements_total"`
	SQLiteWritesTotal  uint64  `json:"sqlite_writes_total"`
	RejectedTotal      uint64  `json:"rejected_total"`
	UptimeSec          int64   `json:"uptime_sec"`
}

func (s *State) Metrics() Metrics {
	g, i, w, r := s.Rates(5)
	return Metrics{Connected: s.ConnectedCount(), GuessesPerSec: g, ImprovementsPerSec: i, SQLiteWritesPerSec: w, RejectedPerSec: r, GuessesTotal: s.guessesTotal.Load(), ImprovementsTotal: s.improvementsTotal.Load(), SQLiteWritesTotal: s.sqliteWritesTotal.Load(), RejectedTotal: s.rejectedTotal.Load(), UptimeSec: int64(time.Since(s.started).Seconds())}
}