// Package bandit is the Phase 13 multi-armed bandit filter tuning engine.
// Operators define filter variants in config; the bandit assigns credit for
// each grab to one variant according to its strategy, building per-variant
// posteriors stored in filter_performance.variant_id.
package bandit

import (
	"context"
	"database/sql"
	"errors"
	"math/rand/v2"
	"sync"
	"time"

	"github.com/rs/zerolog"

	"github.com/operator/command-center/internal/eventbus"
	"github.com/operator/command-center/internal/webhooks"
)

// Strategy is the per-arm selection rule.
type Strategy string

const (
	StrategyThompson    Strategy = "thompson_sampling"
	StrategyEpsilon     Strategy = "epsilon_greedy"
	StrategyUCB1        Strategy = "ucb1"
)

// Variant is one operator-declared candidate.
type Variant struct {
	ID          string `yaml:"id"`
	Description string `yaml:"description,omitempty"`
	// Predicate is a simple expression operating on grab attributes. Phase 13
	// supports a tiny vocabulary: "true", "size_bytes <= N", "group == 'X'",
	// joined by " && " / " || ". The evaluator below understands this subset.
	Predicate string `yaml:"predicate"`
}

// Bandit is one operator-declared bandit definition.
type Bandit struct {
	FilterID  string    `yaml:"filter_id"`
	Variants  []Variant `yaml:"variants"`
	Strategy  Strategy  `yaml:"strategy"`
	Epsilon   float64   `yaml:"epsilon,omitempty"`
	PriorA    int       `yaml:"prior_alpha,omitempty"`
	PriorB    int       `yaml:"prior_beta,omitempty"`
}

// Selector dispatches grab credit. One per process; subscribes to the
// autobrr.grab event topic.
type Selector struct {
	db     *sql.DB
	bus    *eventbus.Bus
	logger zerolog.Logger

	mu      sync.RWMutex
	bandits map[string]Bandit // keyed by FilterID
	rng     *rand.Rand
	wg      sync.WaitGroup
	stop    context.CancelFunc
}

// NewSelector constructs a Selector.
func NewSelector(db *sql.DB, bus *eventbus.Bus, logger zerolog.Logger) *Selector {
	return &Selector{
		db:     db,
		bus:    bus,
		logger: logger.With().Str("component", "bandit").Logger(),
		bandits: map[string]Bandit{},
		rng:    rand.New(rand.NewPCG(uint64(time.Now().UnixNano()), 0)),
	}
}

// Replace swaps the in-memory bandit definitions.
func (s *Selector) Replace(bs []Bandit) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.bandits = map[string]Bandit{}
	for _, b := range bs {
		s.bandits[b.FilterID] = b
	}
}

// Start subscribes to TopicAutobrrGrab and processes events.
func (s *Selector) Start(parent context.Context) {
	ctx, cancel := context.WithCancel(parent)
	s.stop = cancel
	s.wg.Add(1)
	go s.run(ctx)
}

// Stop terminates the subscriber.
func (s *Selector) Stop() {
	if s.stop != nil {
		s.stop()
	}
	s.wg.Wait()
}

func (s *Selector) run(ctx context.Context) {
	defer s.wg.Done()
	sub := s.bus.Subscribe(eventbus.TopicAutobrrGrab)
	defer sub.Close()
	for {
		select {
		case <-ctx.Done():
			return
		case ev, ok := <-sub.Chan():
			if !ok {
				return
			}
			grab, ok := ev.Payload.(webhooks.AutobrrGrabEvent)
			if !ok {
				continue
			}
			s.handleGrab(ctx, grab)
		}
	}
}

// handleGrab selects one variant and updates the existing filter_performance
// row (written by the Phase 5 linker) with the variant_id.
func (s *Selector) handleGrab(ctx context.Context, g webhooks.AutobrrGrabEvent) {
	s.mu.RLock()
	b, ok := s.bandits[g.FilterID]
	s.mu.RUnlock()
	if !ok {
		return
	}
	matched := s.matchingVariants(b, g)
	if len(matched) == 0 {
		return
	}
	picked := s.pick(b, matched)
	if picked == "" {
		return
	}
	// Update the most recent filter_performance row for this filter (the
	// Phase 5 linker just inserted it) with the variant_id. Idempotent: if
	// no row exists yet, this is a no-op.
	_, _ = s.db.ExecContext(ctx, `
		UPDATE filter_performance
		SET variant_id = ?
		WHERE id = (
		  SELECT id FROM filter_performance
		  WHERE filter_external_id = ? AND variant_id IS NULL
		  ORDER BY grabbed_at DESC LIMIT 1
		)
	`, picked, g.FilterID)
}

func (s *Selector) matchingVariants(b Bandit, g webhooks.AutobrrGrabEvent) []Variant {
	out := make([]Variant, 0, len(b.Variants))
	for _, v := range b.Variants {
		if evalPredicate(v.Predicate, g) {
			out = append(out, v)
		}
	}
	return out
}

// pick implements the configured strategy. Phase 13 ships a clean Thompson
// sampling over Beta posteriors built from the filter_performance table.
// Epsilon-greedy + UCB1 are stubs; operators choosing them today get the
// same selection logic as Thompson with a warning.
func (s *Selector) pick(b Bandit, candidates []Variant) string {
	switch b.Strategy {
	case StrategyEpsilon:
		// ε exploration: with probability ε pick a random arm, else the
		// best-mean arm.
		if s.rng.Float64() < b.Epsilon {
			return candidates[s.rng.IntN(len(candidates))].ID
		}
		return s.bestMeanArm(b, candidates)
	case StrategyUCB1:
		return s.ucbArm(b, candidates)
	default:
		return s.thompsonArm(b, candidates)
	}
}

func (s *Selector) thompsonArm(b Bandit, candidates []Variant) string {
	alpha, beta := b.PriorA, b.PriorB
	if alpha < 1 {
		alpha = 1
	}
	if beta < 1 {
		beta = 1
	}
	bestName, bestSample := "", -1.0
	for _, v := range candidates {
		successes, total := s.armCounts(b.FilterID, v.ID)
		a := float64(alpha + successes)
		bb := float64(beta + total - successes)
		sample := s.sampleBeta(a, bb)
		if sample > bestSample {
			bestSample = sample
			bestName = v.ID
		}
	}
	return bestName
}

func (s *Selector) bestMeanArm(b Bandit, candidates []Variant) string {
	bestName, bestMean := candidates[0].ID, -1.0
	for _, v := range candidates {
		successes, total := s.armCounts(b.FilterID, v.ID)
		mean := 0.0
		if total > 0 {
			mean = float64(successes) / float64(total)
		}
		if mean > bestMean {
			bestMean = mean
			bestName = v.ID
		}
	}
	return bestName
}

func (s *Selector) ucbArm(b Bandit, candidates []Variant) string {
	bestName, bestScore := candidates[0].ID, -1.0
	var total int
	armTotals := make(map[string]int, len(candidates))
	armSuccesses := make(map[string]int, len(candidates))
	for _, v := range candidates {
		ss, tt := s.armCounts(b.FilterID, v.ID)
		armSuccesses[v.ID] = ss
		armTotals[v.ID] = tt
		total += tt
	}
	if total == 0 {
		return candidates[0].ID
	}
	for _, v := range candidates {
		ss := armSuccesses[v.ID]
		tt := armTotals[v.ID]
		if tt == 0 {
			return v.ID // explore unseen arm immediately
		}
		mean := float64(ss) / float64(tt)
		bonus := mathSqrt(2 * mathLog(float64(total)) / float64(tt))
		score := mean + bonus
		if score > bestScore {
			bestScore = score
			bestName = v.ID
		}
	}
	return bestName
}

// armCounts returns (successes, total) where success = final_uploaded_bytes > size_threshold;
// for Phase 13 we approximate "success" as any grab that produced final_uploaded_bytes
// at all, so we get a binomial over "did this grab pan out".
func (s *Selector) armCounts(filterID, variantID string) (int, int) {
	var total, succ int
	_ = s.db.QueryRow(`
		SELECT COUNT(*),
		       COUNT(CASE WHEN final_uploaded_bytes IS NOT NULL AND final_uploaded_bytes > 0 THEN 1 END)
		FROM filter_performance
		WHERE filter_external_id = ? AND variant_id = ?
	`, filterID, variantID).Scan(&total, &succ)
	return succ, total
}

// sampleBeta uses the ratio-of-Gammas method via NormFloat64; for small α/β
// at single-operator scale this is plenty fast and accurate.
func (s *Selector) sampleBeta(alpha, beta float64) float64 {
	x := sampleGamma(s.rng, alpha)
	y := sampleGamma(s.rng, beta)
	if x+y == 0 {
		return 0.5
	}
	return x / (x + y)
}

// sampleGamma uses Marsaglia & Tsang's method for shape ≥ 1, with the
// integer-shape special case for shape < 1.
func sampleGamma(rng *rand.Rand, shape float64) float64 {
	if shape < 1 {
		return sampleGamma(rng, shape+1) * mathPow(rng.Float64(), 1/shape)
	}
	d := shape - 1.0/3.0
	c := 1.0 / mathSqrt(9.0*d)
	for {
		x := rng.NormFloat64()
		v := 1.0 + c*x
		if v <= 0 {
			continue
		}
		v = v * v * v
		u := rng.Float64()
		if u < 1-0.0331*x*x*x*x {
			return d * v
		}
		if mathLog(u) < 0.5*x*x+d*(1-v+mathLog(v)) {
			return d * v
		}
	}
}

var _ = errors.New