MOPAC/internal/quota/decision.go

package quota

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"

	"ukrrs.com/mopac/harness/internal/config"
)

// Decision is the back-pressure verdict for one candidate dispatch.
type Decision struct {
	Action string // "allow" | "defer"
	Reason string // human/log line; empty on allow unless Peak is set
	Peak   bool   // decision was made inside the peak window (context)
}

// Action constants.
const (
	ActionAllow = "allow"
	ActionDefer = "defer"
)

func allow(peak bool) Decision { return Decision{Action: ActionAllow, Peak: peak} }

func defer_(reason string, peak bool) Decision {
	return Decision{Action: ActionDefer, Reason: reason, Peak: peak}
}

// Gate is the loop's back-pressure consultant: quota snapshot (polled or
// estimated), peak schedule, thresholds. It never blocks startup on redis
// or the provider being down; it degrades to local estimates and allows.
type Gate struct {
	cfg    config.QuotaConfig
	sched  Schedule
	shared *SharedState
	keys   *config.KeyResolver
	http   *http.Client
	now    func() time.Time

	mu         sync.Mutex
	lastSnap   *QuotaSnapshot // last successful poll (cache across failures)
	local5h    float64        // process-local estimate fallbacks (redis down)
	localWeek  float64
	win5hID    string
	winWeekID  string
}

// NewGate builds the gate from the [quota] config. shared may be nil
// (local-only estimates); keys may be nil when usage_url is unset.
func NewGate(cfg config.QuotaConfig, keys *config.KeyResolver, shared *SharedState) (*Gate, error) {
	sched, err := NewSchedule(cfg.PeakStart, cfg.PeakEnd, cfg.Timezone, cfg.PeakWeekdaysOnly)
	if err != nil {
		return nil, fmt.Errorf("quota schedule: %w", err)
	}
	return &Gate{
		cfg:    cfg,
		sched:  sched,
		shared: shared,
		keys:   keys,
		http:   &http.Client{Timeout: 15 * time.Second},
		now:    time.Now,
	}, nil
}

// SetClock overrides the gate's clock (tests inject a fake).
func (g *Gate) SetClock(now func() time.Time) { g.now = now }

// InPeak exposes the schedule verdict at the gate's clock.
func (g *Gate) InPeak() bool { return g.sched.InPeak(g.now()) }

// Schedule returns the parsed peak window (status output).
func (g *Gate) Schedule() Schedule { return g.sched }

// Snapshot returns the freshest quota state: a live poll when the provider
// endpoint is configured and the cache is stale, else the redis-published
// snapshot from any instance, else the local cache, else the estimate
// synthesized from this instance's recorded consumption. The returned
// snapshot is never nil when the gate is enabled.
func (g *Gate) Snapshot(ctx context.Context) *QuotaSnapshot {
	now := g.now()
	if g.cfg.UsageURL != "" {
		if s := g.cachedPoll(); s != nil && now.Sub(s.FetchedAt) < time.Duration(g.cfg.PollIntervalSecs)*time.Second {
			return s
		}
		if s, err := g.poll(ctx, now); err == nil {
			return s
		}
	}
	if g.shared != nil {
		if s, err := g.shared.LoadSnapshot(g.cfg.Account); err == nil && s != nil &&
			now.Sub(s.FetchedAt) < 2*time.Duration(g.cfg.PollIntervalSecs)*time.Second {
			return s
		}
	}
	if s := g.cachedPoll(); s != nil {
		return s
	}
	return g.estimateSnapshot(now)
}

func (g *Gate) cachedPoll() *QuotaSnapshot {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.lastSnap
}

func (g *Gate) poll(ctx context.Context, now time.Time) (*QuotaSnapshot, error) {
	key, err := g.keys.Resolve(ctx, g.cfg.KeyRef)
	if err != nil {
		return nil, fmt.Errorf("quota key: %w", err)
	}
	body, err := FetchUsage(ctx, g.http, g.cfg.UsageURL, key)
	if err != nil {
		return nil, err
	}
	snap, err := ParseUsage(g.cfg.Account, now, body)
	if err != nil {
		return nil, err
	}
	g.mu.Lock()
	g.lastSnap = snap
	g.mu.Unlock()
	_ = g.shared.PublishSnapshot(snap, time.Duration(g.cfg.PollIntervalSecs)*3*time.Second)
	return snap, nil
}

// window ids: the 5h bucket advances in fixed 5h steps from the epoch; the
// weekly bucket is the Monday 00:00 of the schedule's TZ (z.ai weekly
// credits reset 7 days after activation — an approximation until the usage
// endpoint ships reset_at, noted in the REPORT).
func (g *Gate) windowIDs(now time.Time) (w5h, week string) {
	five := 5 * time.Hour
	w5h = fmt.Sprintf("%d", now.Unix()/int64(five.Seconds()))
	mon := now.In(g.sched.Loc)
	for mon.Weekday() != time.Monday {
		mon = mon.AddDate(0, 0, -1)
	}
	y, m, d := mon.Date()
	week = fmt.Sprintf("%04d%02d%02d", y, int(m), d)
	return w5h, week
}

// RecordTurn records one completed turn's estimated credits into the shared
// state (redis when configured; process-local maps always, so a down redis
// never loses this instance's own accounting).
func (g *Gate) RecordTurn(in EstimateTurnInput) float64 {
	credits := EstimateCredits(in.Model, in.PromptTokens, in.CachedTokens, in.CompletionTokens, in.Peak)
	now := g.now()
	w5h, week := g.windowIDs(now)

	g.mu.Lock()
	if g.win5hID != w5h { // 5h window rolled: reset the local counter
		g.local5h, g.win5hID = 0, w5h
	}
	if g.winWeekID != week {
		g.localWeek, g.winWeekID = 0, week
	}
	g.local5h += credits
	g.localWeek += credits
	g.mu.Unlock()

	if g.shared != nil {
		k5 := fmt.Sprintf("mopac:quota:%s:est:5h:%s", g.cfg.Account, w5h)
		kw := fmt.Sprintf("mopac:quota:%s:est:weekly:%s", g.cfg.Account, week)
		if _, err := g.shared.incrByFloat(k5, credits); err == nil {
			_, _ = g.shared.command("EXPIRE", k5, "21600") // 5h + 1h slack
		}
		_, _ = g.shared.incrByFloat(kw, credits)
	}
	return credits
}

// estimateSnapshot synthesizes quota state from recorded consumption against
// the configured plan limits (the no-endpoint mode).
func (g *Gate) estimateSnapshot(now time.Time) *QuotaSnapshot {
	w5h, week := g.windowIDs(now)
	var used5, usedW float64
	if g.shared != nil {
		if v, err := g.shared.get(fmt.Sprintf("mopac:quota:%s:est:5h:%s", g.cfg.Account, w5h)); err == nil && v != "" {
			fmt.Sscanf(v, "%g", &used5)
		}
		if v, err := g.shared.get(fmt.Sprintf("mopac:quota:%s:est:weekly:%s", g.cfg.Account, week)); err == nil && v != "" {
			fmt.Sscanf(v, "%g", &usedW)
		}
	}
	if used5 == 0 || usedW == 0 {
		g.mu.Lock()
		if g.win5hID == w5h {
			used5 = max(used5, g.local5h)
		}
		if g.winWeekID == week {
			usedW = max(usedW, g.localWeek)
		}
		g.mu.Unlock()
	}
	src := "estimate"
	if g.shared != nil {
		src = "estimate+redis"
	}
	return &QuotaSnapshot{
		Account:   g.cfg.Account,
		Source:    src,
		FetchedAt: now,
		Buckets: []Bucket{
			{ID: Bucket5h, Used: used5, Limit: g.cfg.Plan5hCredits,
				WindowReset: now.Add(5 * time.Hour)},
			{ID: BucketWeekly, Used: usedW, Limit: g.cfg.PlanWeeklyCredits,
				WindowReset: nextWeekStart(now, g.sched.Loc)},
		},
	}
}

func nextWeekStart(now time.Time, loc *time.Location) time.Time {
	t := now.In(loc)
	for t.Weekday() != time.Monday {
		t = t.AddDate(0, 0, 1)
	}
	y, m, d := t.Date()
	return time.Date(y, m, d, 0, 0, 0, 0, loc).UTC()
}

// Decide is the pre-dispatch consultation. Priority: hard quota wall first
// (the 19:00 failure mode), then peak-window class restriction, then
// soft-quota heavy-class deferral. Unknown quota state never defers.
func (g *Gate) Decide(ctx context.Context, class string) Decision {
	peak := g.InPeak()
	snap := g.Snapshot(ctx)
	ratio := snap.MaxUsedPct()

	if ratio >= g.cfg.BlockAtPct {
		return defer_(fmt.Sprintf("quota: %s bucket at %.0f%% (>= block %.0f%%): all classes deferred until reset (source %s)",
			worstBucketID(snap), ratio, g.cfg.BlockAtPct, snap.Source), peak)
	}
	if peak && !g.peakClass(class) {
		return defer_(fmt.Sprintf("peak window %s-%s %s: class %q deferred to off-peak (flash-tier classes only: %s)",
			g.cfg.PeakStart, g.cfg.PeakEnd, g.cfg.Timezone, class, strings.Join(g.cfg.PeakClasses, ",")), peak)
	}
	if ratio >= g.cfg.DeferAtPct && !g.peakClass(class) {
		return defer_(fmt.Sprintf("quota: %s bucket at %.0f%% (>= defer %.0f%%): heavy class %q deferred; LLM-lite continues (source %s)",
			worstBucketID(snap), ratio, g.cfg.DeferAtPct, class, snap.Source), peak)
	}
	return allow(peak)
}

func (g *Gate) peakClass(class string) bool {
	for _, c := range g.cfg.PeakClasses {
		if c == class {
			return true
		}
	}
	return false
}

func worstBucketID(s *QuotaSnapshot) string {
	id, max := "", 0.0
	for _, b := range s.Buckets {
		if p := b.UsedPct(); p > max {
			max, id = p, b.ID
		}
	}
	return id
}

// StatusLine renders the one-line quota status for loop startup / `harness
// quota status`.
func (g *Gate) StatusLine(snap *QuotaSnapshot) string {
	parts := make([]string, 0, len(snap.Buckets)+2)
	parts = append(parts, fmt.Sprintf("account=%s source=%s", snap.Account, snap.Source))
	for _, b := range snap.Buckets {
		parts = append(parts, fmt.Sprintf("%s=%.0f/%.0f(%.0f%%)", b.ID, b.Used, b.Limit, b.UsedPct()))
	}
	if g.InPeak() {
		parts = append(parts, "PEAK")
	}
	return strings.Join(parts, " ")
}

// SortedClasses returns PeakClasses sorted (stable status output).
func (g *Gate) SortedClasses() []string {
	out := append([]string{}, g.cfg.PeakClasses...)
	sort.Strings(out)
	return out
}