AetherForge/server/internal/api/ws_beacon_integration_test.go

package api

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"

	"github.com/gorilla/websocket"
)

// TestIntegrationAuthStatsTickStatsBatchEndToEnd verifies auth → stats tick →
// coalesced stats_batch delivery to a dashboard WebSocket client.
func TestIntegrationAuthStatsTickStatsBatchEndToEnd(t *testing.T) {
	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(dashSrv.Close)
	dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
	if err != nil {
		t.Fatalf("dial dashboard: %v", err)
	}
	t.Cleanup(func() { _ = dashConn.Close() })

	type batchResult struct {
		updates []map[string]interface{}
		err     string
	}
	batchCh := make(chan batchResult, 1)
	go func() {
		_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := dashConn.ReadJSON(&msg); err != nil {
				batchCh <- batchResult{err: err.Error()}
				return
			}
			if msg.Type != "stats_batch" {
				continue
			}
			var body struct {
				Updates []json.RawMessage `json:"updates"`
			}
			if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
				batchCh <- batchResult{err: parseErr.Error()}
				return
			}
			updates := make([]map[string]interface{}, 0, len(body.Updates))
			for _, raw := range body.Updates {
				var u map[string]interface{}
				if json.Unmarshal(raw, &u) != nil {
					continue
				}
				updates = append(updates, u)
			}
			batchCh <- batchResult{updates: updates}
			return
		}
	}()

	agentID := "auth-stats-agent"
	conn := connectTestAgent(t, hub, agentID)

	statsPayload, _ := json.Marshal(map[string]interface{}{
		"hashrate_15s":     42.0,
		"hashrate_1m":      40.0,
		"hashrate_15m":     38.0,
		"shares_submitted": 3,
		"shares_accepted":  2,
		"cpu_usage_pct":    11.0,
		"memory_usage_pct": 22.0,
		"uptime_seconds":   120,
		"mining_hashrate":  42.0,
		"lotl_tier":        "cpu_inprocess",
		"parent_agent_id":  "parent-abc",
		"spread_generation": 1,
		"spread_strain":    "#112233",
	})
	if err := conn.WriteJSON(Message{Type: "stats", Payload: statsPayload}); err != nil {
		t.Fatal(err)
	}
	stopStatsBatchTimer(hub)
	hub.flushStatsBatch()

	select {
	case r := <-batchCh:
		if r.err != "" {
			t.Fatalf("dashboard did not receive stats_batch: %s", r.err)
		}
		if len(r.updates) != 1 {
			t.Fatalf("expected 1 update, got %d: %+v", len(r.updates), r.updates)
		}
		u := r.updates[0]
		if u["agent_id"] != agentID {
			t.Errorf("agent_id = %v", u["agent_id"])
		}
		if u["hashrate_15s"] != 42.0 {
			t.Errorf("hashrate_15s = %v", u["hashrate_15s"])
		}
		if u["mining_hashrate"] != 42.0 {
			t.Errorf("mining_hashrate = %v", u["mining_hashrate"])
		}
		if u["lotl_tier"] != "cpu_inprocess" {
			t.Errorf("lotl_tier = %v", u["lotl_tier"])
		}
		if u["parent_agent_id"] != "parent-abc" {
			t.Errorf("parent_agent_id = %v", u["parent_agent_id"])
		}
		if u["spread_generation"] != float64(1) {
			t.Errorf("spread_generation = %v", u["spread_generation"])
		}
		if u["spread_strain"] != "#112233" {
			t.Errorf("spread_strain = %v", u["spread_strain"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for stats_batch after auth+stats tick")
	}

	agent, err := database.GetAgent(agentID)
	if err != nil {
		t.Fatal(err)
	}
	if agent.Status != "online" {
		t.Errorf("agent status = %q, want online", agent.Status)
	}
	if agent.Hashrate15s != 42.0 {
		t.Errorf("db hashrate_15s = %v", agent.Hashrate15s)
	}
}

// TestIntegrationBeaconRegistrationHeartbeatLifecycle covers HTTPS beacon
// registration, heartbeat reachability, queued command delivery, and result relay.
func TestIntegrationBeaconRegistrationHeartbeatLifecycle(t *testing.T) {
	resetAuthState(t)
	const secret = "beacon-lifecycle-secret"
	SetAgentPathSecret(secret)

	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetFleetSecret(secret)

	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(dashSrv.Close)
	dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
	if err != nil {
		t.Fatalf("dial dashboard: %v", err)
	}
	t.Cleanup(func() { _ = dashConn.Close() })

	beaconHandler := basicAuthMiddleware(http.HandlerFunc(hub.HandleAgentBeacon))
	beaconResultHandler := basicAuthMiddleware(http.HandlerFunc(hub.HandleAgentBeaconResult))

	postBeacon := func(agentID, hostname string, hashrate float64) *httptest.ResponseRecorder {
		t.Helper()
		body, _ := json.Marshal(map[string]interface{}{
			"agent_id": agentID,
			"hostname": hostname,
			"version":  "1.0",
			"stats": map[string]interface{}{
				"hashrate_15s": hashrate,
				"hashrate_1m":  hashrate,
				"hashrate_15m": hashrate,
			},
		})
		req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/beacon", bytes.NewReader(body))
		req.Header.Set("X-Fleet-Secret", secret)
		rec := httptest.NewRecorder()
		beaconHandler.ServeHTTP(rec, req)
		return rec
	}

	agentID := "beacon-new-agent"
	rec := postBeacon(agentID, "BEACON-HOST", 55.0)
	if rec.Code != http.StatusOK {
		t.Fatalf("first beacon: %d %s", rec.Code, rec.Body.String())
	}
	agent, err := database.GetAgent(agentID)
	if err != nil {
		t.Fatal(err)
	}
	if agent.Name != "BEACON-HOST" {
		t.Errorf("registered name = %q, want BEACON-HOST", agent.Name)
	}
	if !hub.isAgentBeaconReachable(agentID) {
		t.Fatal("agent should be beacon-reachable after first heartbeat")
	}

	rec = postBeacon(agentID, "BEACON-HOST", 60.0)
	if rec.Code != http.StatusOK {
		t.Fatalf("second beacon heartbeat: %d", rec.Code)
	}
	if !hub.EnqueueBeaconCommand(agentID, "pause", nil) {
		t.Fatal("enqueue pause should succeed while beacon reachable")
	}

	rec = postBeacon(agentID, "BEACON-HOST", 65.0)
	if rec.Code != http.StatusOK {
		t.Fatalf("beacon with commands: %d", rec.Code)
	}
	var beaconResp beaconResponse
	if err := json.Unmarshal(rec.Body.Bytes(), &beaconResp); err != nil {
		t.Fatal(err)
	}
	if len(beaconResp.Commands) != 1 || beaconResp.Commands[0].Action != "pause" {
		t.Fatalf("expected pause command, got %+v", beaconResp.Commands)
	}

	type cmdResult struct {
		body map[string]interface{}
		err  string
	}
	resultCh := make(chan cmdResult, 1)
	go func() {
		_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := dashConn.ReadJSON(&msg); err != nil {
				resultCh <- cmdResult{err: err.Error()}
				return
			}
			if msg.Type != "command_result" {
				continue
			}
			var body map[string]interface{}
			if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
				resultCh <- cmdResult{err: parseErr.Error()}
				return
			}
			if body["transport"] != "https_beacon" {
				continue
			}
			resultCh <- cmdResult{body: body}
			return
		}
	}()

	resultBody, _ := json.Marshal(map[string]interface{}{
		"agent_id": agentID,
		"action":   "pause",
		"success":  true,
		"message":  "paused via beacon",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/beacon/result", bytes.NewReader(resultBody))
	req.Header.Set("X-Fleet-Secret", secret)
	rec = httptest.NewRecorder()
	beaconResultHandler.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("beacon result: %d %s", rec.Code, rec.Body.String())
	}

	select {
	case r := <-resultCh:
		if r.err != "" {
			t.Fatalf("dashboard did not receive beacon command_result: %s", r.err)
		}
		if r.body["agent_id"] != agentID {
			t.Errorf("agent_id = %v", r.body["agent_id"])
		}
		if r.body["action"] != "pause" {
			t.Errorf("action = %v", r.body["action"])
		}
		if r.body["message"] != "paused via beacon" {
			t.Errorf("message = %v", r.body["message"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for beacon command_result broadcast")
	}
}

// TestIntegrationBeaconClearsOnWSReconnect verifies beacon transport state is
// cleared when the agent reconnects over WebSocket.
func TestIntegrationBeaconClearsOnWSReconnect(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	_ = database.UpsertAgent(&models.Agent{ID: "beacon-ws-agent", Name: "host", Status: "offline"})

	hub := NewWSHub(database)
	hub.MarkBeaconSeen("beacon-ws-agent")
	_ = hub.EnqueueBeaconCommand("beacon-ws-agent", "resume", nil)
	if !hub.isAgentBeaconReachable("beacon-ws-agent") {
		t.Fatal("expected beacon reachable before WS auth")
	}

	connectTestAgent(t, hub, "beacon-ws-agent")

	if hub.isAgentBeaconReachable("beacon-ws-agent") {
		t.Fatal("beacon state should be cleared after WS reconnect")
	}
	if len(hub.dequeueBeaconCommands("beacon-ws-agent")) != 0 {
		t.Fatal("beacon command queue should be empty after WS reconnect")
	}
}

// TestIntegrationCommandDispatchExecShellRoundTrip sends exec_shell over WS and
// verifies the agent receives the framed command payload.
func TestIntegrationCommandDispatchExecShellRoundTrip(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "exec-shell-agent"
	conn := connectTestAgent(t, hub, agentID)

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

	const shellCmd = "echo integration-exec"
	if err := hub.SendAgentCommand(agentID, "exec_shell", map[string]interface{}{
		"command": shellCmd,
	}); err != nil {
		t.Fatal(err)
	}

	select {
	case r := <-cmdCh:
		if r.err != "" {
			t.Fatalf("agent did not receive command: %s", r.err)
		}
		var payload map[string]interface{}
		if err := json.Unmarshal(r.cmd.Payload, &payload); err != nil {
			t.Fatal(err)
		}
		if payload["action"] != "exec_shell" {
			t.Errorf("action = %v", payload["action"])
		}
		if payload["command"] != shellCmd {
			t.Errorf("command = %v", payload["command"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for exec_shell command")
	}
}

// TestIntegrationCommandDispatchMiningDiagnosticsRoundTrip sends
// mining_diagnostics and verifies the agent receives it.
func TestIntegrationCommandDispatchMiningDiagnosticsRoundTrip(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "mining-diag-agent"
	conn := connectTestAgent(t, hub, agentID)

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

	if err := hub.SendAgentCommand(agentID, "mining_diagnostics", nil); err != nil {
		t.Fatal(err)
	}

	select {
	case r := <-cmdCh:
		if r.err != "" {
			t.Fatalf("agent did not receive command: %s", r.err)
		}
		var payload map[string]interface{}
		if err := json.Unmarshal(r.cmd.Payload, &payload); err != nil {
			t.Fatal(err)
		}
		if payload["action"] != "mining_diagnostics" {
			t.Errorf("action = %v", payload["action"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for mining_diagnostics command")
	}
}

// TestIntegrationAISnapshotRequestFlow sends ai_snapshot_request over WS and
// verifies an ai_snapshot reply is cached in hub telemetry for Fleet AI.
func TestIntegrationAISnapshotRequestFlow(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "ai-snapshot-agent"
	conn := connectIntelAgent(t, hub, agentID, nil)

	pushStuckAgentTelemetry(t, conn)

	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		hub.mu.RLock()
		tel := hub.agentLiveTelemetry[agentID]
		hub.mu.RUnlock()
		if tel != nil {
			if stuck, _ := tel["stuck"].(bool); stuck {
				return
			}
		}
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatal("ai_snapshot telemetry never cached in hub")
}

// TestIntegrationAgentDisconnectCleanup verifies WS disconnect clears live hub
// state, marks the agent offline, and broadcasts agent_offline to dashboards.
func TestIntegrationAgentDisconnectCleanup(t *testing.T) {
	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	dashSrv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(dashSrv.Close)
	dashURL := "ws" + strings.TrimPrefix(dashSrv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	dashConn, _, err := websocket.DefaultDialer.Dial(dashURL, nil)
	if err != nil {
		t.Fatalf("dial dashboard: %v", err)
	}
	t.Cleanup(func() { _ = dashConn.Close() })

	agentID := "disconnect-cleanup-agent"
	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)
	}

	authAgentConn(t, conn, map[string]interface{}{
		"agent_id": agentID,
		"hostname": "cleanup-host",
		"version":  "1.0",
	})
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if hub.isAgentConnected(agentID) {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if !hub.isAgentConnected(agentID) {
		t.Fatal("agent should be connected after auth")
	}

	logPayload, _ := json.Marshal(map[string]interface{}{"content": "tail-line", "lines": 1})
	if err := conn.WriteJSON(Message{Type: "log_tail", Payload: logPayload}); err != nil {
		t.Fatal(err)
	}
	time.Sleep(30 * time.Millisecond)
	if got := hub.GetAgentLog(agentID); got != "tail-line" {
		t.Fatalf("log tail = %q", got)
	}

	offlineCh := make(chan map[string]interface{}, 1)
	go func() {
		_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := dashConn.ReadJSON(&msg); err != nil {
				return
			}
			if msg.Type != "agent_offline" {
				continue
			}
			var body map[string]interface{}
			if json.Unmarshal(msg.Payload, &body) != nil {
				continue
			}
			if body["agent_id"] == agentID {
				offlineCh <- body
				return
			}
		}
	}()

	_ = conn.Close()

	waitDeadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(waitDeadline) {
		if !hub.isAgentConnected(agentID) {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if hub.isAgentConnected(agentID) {
		t.Fatal("agent should be disconnected after conn close")
	}
	if hub.GetAgentLog(agentID) != "" {
		t.Fatal("agent log cache should be cleared on disconnect")
	}

	select {
	case body := <-offlineCh:
		if body["agent_id"] != agentID {
			t.Errorf("offline agent_id = %v", body["agent_id"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for agent_offline broadcast")
	}

	agent, err := database.GetAgent(agentID)
	if err != nil {
		t.Fatal(err)
	}
	if agent.Status != "offline" {
		t.Errorf("db status = %q, want offline", agent.Status)
	}
}