AetherForge/server/internal/api/websocket_test.go

package api

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

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

	"github.com/gorilla/websocket"
)

const wsDefaultPingIntervalSec = 30

func resetWSAuthUsers(t *testing.T, user, pass string) {
	t.Helper()
	hashed, err := hashPassword(pass)
	if err != nil {
		t.Fatal(err)
	}
	usersMu.Lock()
	authUsers = map[string]string{user: hashed}
	usersMu.Unlock()
	t.Cleanup(func() {
		usersMu.Lock()
		authUsers = map[string]string{}
		usersMu.Unlock()
	})
}

func wsDashboardToken(user, pass string) string {
	return base64.StdEncoding.EncodeToString([]byte(user + ":" + pass))
}

func dialAgentWS(t *testing.T, hub *WSHub) (*websocket.Conn, string) {
	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: %v", err)
	}
	t.Cleanup(func() { _ = conn.Close() })
	return conn, wsURL
}

func authAgentConn(t *testing.T, conn *websocket.Conn, payload map[string]interface{}) Message {
	t.Helper()
	data, _ := json.Marshal(payload)
	if err := conn.WriteJSON(Message{Type: "auth", Payload: data}); err != nil {
		t.Fatal(err)
	}
	var resp Message
	if err := conn.ReadJSON(&resp); err != nil {
		t.Fatalf("read auth_response: %v", err)
	}
	return resp
}

func TestWSHubPingIntervalConstants(t *testing.T) {
	hub := NewWSHub(nil)
	if hub.pingIntervalSec != wsDefaultPingIntervalSec {
		t.Fatalf("default ping interval = %d", hub.pingIntervalSec)
	}
	hub.SetPingInterval(5)
	if hub.pingIntervalSec != wsDefaultPingIntervalSec {
		t.Fatalf("below-minimum ping should clamp to %d, got %d", wsDefaultPingIntervalSec, hub.pingIntervalSec)
	}
	hub.SetPingInterval(15)
	if hub.pingIntervalSec != 15 {
		t.Fatalf("expected 15, got %d", hub.pingIntervalSec)
	}
	if hub.pingInterval().Seconds() != 15 {
		t.Fatalf("pingInterval duration = %v", hub.pingInterval())
	}
}

func TestCheckDashboardWSTokenBcryptUser(t *testing.T) {
	resetWSAuthUsers(t, "dash", "secret-pass")
	req := httptest.NewRequest(http.MethodGet, "/ws/dashboard?token="+wsDashboardToken("dash", "secret-pass"), nil)
	if !checkDashboardWSToken(req) {
		t.Fatal("valid bcrypt user token should pass")
	}
	req = httptest.NewRequest(http.MethodGet, "/ws/dashboard?token="+wsDashboardToken("dash", "wrong"), nil)
	if checkDashboardWSToken(req) {
		t.Fatal("wrong password should fail")
	}
	req = httptest.NewRequest(http.MethodGet, "/ws/dashboard", nil)
	if checkDashboardWSToken(req) {
		t.Fatal("missing token should fail")
	}
}

func TestCheckDashboardWSTicketOneTime(t *testing.T) {
	ticket := issueWSTicket("dash")
	req := httptest.NewRequest(http.MethodGet, "/ws/dashboard?ticket="+ticket, nil)
	if !checkDashboardWSToken(req) {
		t.Fatal("valid ticket should pass")
	}
	req = httptest.NewRequest(http.MethodGet, "/ws/dashboard?ticket="+ticket, nil)
	if checkDashboardWSToken(req) {
		t.Fatal("ticket should be one-time use")
	}
}

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

	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(srv.Close)

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
	_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err == nil {
		t.Fatal("expected dial failure without token")
	}
	if resp == nil || resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("expected 401 upgrade rejection, got err=%v status=%v", err, resp)
	}
}

func TestHandleDashboardWSAuthorizedInit(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)

	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(srv.Close)

	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "?token=" + wsDashboardToken(testAuthUser, testAuthPass)
	conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("dial: %v status=%v", err, resp)
	}
	t.Cleanup(func() { _ = conn.Close() })

	var msg Message
	if err := conn.ReadJSON(&msg); err != nil {
		t.Fatalf("read init: %v", err)
	}
	if msg.Type != "init" {
		t.Fatalf("expected init message, got %q", msg.Type)
	}
}

func resetWSAuthUsersMulti(t *testing.T, creds map[string]string) {
	t.Helper()
	users := make(map[string]string, len(creds))
	for user, pass := range creds {
		hashed, err := hashPassword(pass)
		if err != nil {
			t.Fatal(err)
		}
		users[user] = hashed
	}
	usersMu.Lock()
	authUsers = users
	usersMu.Unlock()
	t.Cleanup(func() {
		usersMu.Lock()
		authUsers = map[string]string{}
		usersMu.Unlock()
	})
}

func TestHandleDashboardWSPresence(t *testing.T) {
	resetWSAuthUsersMulti(t, map[string]string{"india": "secret-pass", "comrade": "secret-pass"})
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	srv := httptest.NewServer(http.HandlerFunc(hub.HandleDashboardWS))
	t.Cleanup(srv.Close)
	base := "ws" + strings.TrimPrefix(srv.URL, "http")

	dial := func(user string) *websocket.Conn {
		t.Helper()
		conn, _, err := websocket.DefaultDialer.Dial(base+"?token="+wsDashboardToken(user, "secret-pass"), nil)
		if err != nil {
			t.Fatalf("dial %s: %v", user, err)
		}
		t.Cleanup(func() { _ = conn.Close() })
		var init Message
		if err := conn.ReadJSON(&init); err != nil || init.Type != "init" {
			t.Fatalf("read init for %s: %v type=%q", user, err, init.Type)
		}
		var snap Message
		if err := conn.ReadJSON(&snap); err != nil || snap.Type != "presence_snapshot" {
			t.Fatalf("read presence_snapshot for %s: %v type=%q", user, err, snap.Type)
		}
		return conn
	}

	connA := dial("india")
	connB := dial("comrade")

	waitForMessage := func(conn *websocket.Conn, wantType, wantUser string, check func(map[string]interface{}) bool) {
		t.Helper()
		deadline := time.Now().Add(3 * time.Second)
		for time.Now().Before(deadline) {
			_ = conn.SetReadDeadline(time.Now().Add(250 * time.Millisecond))
			var msg Message
			if err := conn.ReadJSON(&msg); err != nil {
				continue
			}
			if msg.Type != wantType {
				continue
			}
			var body map[string]interface{}
			if err := json.Unmarshal(msg.Payload, &body); err != nil {
				continue
			}
			if wantUser != "" && body["user"] != wantUser {
				continue
			}
			if check != nil && !check(body) {
				continue
			}
			return
		}
		t.Fatalf("timed out waiting for %s user=%q", wantType, wantUser)
	}

	if err := connA.WriteJSON(Message{Type: "presence_page", Payload: mustMarshal(map[string]string{"page": "/crucible"})}); err != nil {
		t.Fatal(err)
	}
	waitForMessage(connB, "presence_update", "india", func(body map[string]interface{}) bool {
		return body["page"] == "/crucible" && body["online"] == true
	})

	if err := connA.WriteJSON(Message{Type: "notes_typing", Payload: mustMarshal(map[string]bool{"active": true})}); err != nil {
		t.Fatal(err)
	}
	waitForMessage(connB, "notes_typing", "india", func(body map[string]interface{}) bool {
		return body["active"] == true
	})
}

func TestHandleAgentWSBadFleetSecret(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetFleetSecret("required-secret")

	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "agent-bad-secret", "fleet_secret": "wrong", "hostname": "host",
	})
	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if body["success"] != false {
		t.Fatalf("expected auth failure, got %+v", body)
	}
	if hub.isAgentConnected("agent-bad-secret") {
		t.Fatal("agent should not register with bad fleet secret")
	}
}

func TestHandleAgentWSStatsAndLogTail(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "stats-agent"
	conn := connectTestAgent(t, hub, agentID)

	statsPayload, _ := json.Marshal(map[string]interface{}{
		"hashrate_15s": 100.0, "hashrate_1m": 90.0, "hashrate_15m": 80.0,
		"shares_submitted": 5, "shares_accepted": 4,
		"cpu_usage_pct": 12.5, "memory_usage_pct": 40.0, "uptime_seconds": 60,
	})
	if err := conn.WriteJSON(Message{Type: "stats", Payload: statsPayload}); err != nil {
		t.Fatal(err)
	}

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

	cmdPayload, _ := json.Marshal(map[string]interface{}{"action": "exec", "success": true})
	if err := conn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
		t.Fatal(err)
	}
}

func TestHandleAgentWSMaxAgentsPolicy(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	hub.SetServerPolicy(ServerPolicy{MaxAgents: 1})

	conn1, _ := dialAgentWS(t, hub)
	authAgentConn(t, conn1, map[string]interface{}{"agent_id": "first", "hostname": "h1"})

	conn2, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn2, map[string]interface{}{"agent_id": "second", "hostname": "h2"})
	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if body["success"] != false {
		t.Fatalf("second agent should be rejected at max=1: %+v", body)
	}
}

func TestHandleAgentWSInvalidAuthPayload(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	conn, _ := dialAgentWS(t, hub)
	if err := conn.WriteJSON(Message{Type: "auth", Payload: json.RawMessage(`"not-an-object"`)}); err != nil {
		t.Fatal(err)
	}
	var resp Message
	if err := conn.ReadJSON(&resp); err != nil {
		t.Fatal(err)
	}
	var body map[string]interface{}
	_ = json.Unmarshal(resp.Payload, &body)
	if body["success"] != false {
		t.Fatalf("invalid auth payload should fail: %+v", body)
	}
}

func TestWSHubSendAgentCommandNotConnected(t *testing.T) {
	hub := NewWSHub(nil)
	if err := hub.SendAgentCommand("missing", "restart", nil); err == nil {
		t.Fatal("expected error for disconnected agent")
	}
}

func TestWSHubBroadcastHelpers(t *testing.T) {
	hub := NewWSHub(nil)
	hub.BroadcastServerLog("   ")
	hub.BroadcastFleetAlert(map[string]string{"level": "info"})
	hub.BroadcastPoolStatus(map[string]string{"connected": "true"})
	hub.BroadcastAIActivity(map[string]string{"agent_id": "a"})
}

func TestWSHubEnrichAgentsCapabilities(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	agentID := "cap-agent"
	connectTestAgent(t, hub, agentID)

	agents := []*models.Agent{{ID: agentID, Name: "x"}}
	hub.enrichAgentsCapabilities(agents)
	if agents[0].Capabilities == nil {
		t.Fatal("expected capabilities enrichment")
	}
}

func TestWSHubAgentPoolConfigDefaults(t *testing.T) {
	hub := NewWSHub(nil)
	hub.defaultPool = pool.Config{Host: "primary.pool", Port: 3333, Wallet: "48wallet", Password: "pw"}
	hub.agentConfigs["a1"] = AgentForgeConfig{PoolHost: "custom.pool", PoolPort: 4444}
	cfg := hub.agentPoolConfig("a1")
	if cfg.Host != "custom.pool" || cfg.Port != 4444 {
		t.Fatalf("unexpected pool cfg: %+v", cfg)
	}
	cfg = hub.agentPoolConfig("missing")
	if cfg.Host != "primary.pool" {
		t.Fatalf("missing agent should use default pool: %+v", cfg)
	}
}

func TestWSHubConnectedAgentCount(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)
	if hub.connectedAgentCount() != 0 {
		t.Fatal("expected zero agents initially")
	}
	connectTestAgent(t, hub, "count-agent")
	if hub.connectedAgentCount() != 1 {
		t.Fatalf("expected 1 connected agent, got %d", hub.connectedAgentCount())
	}
}

// TestAgentNamePreservedOnReconnect checks that an operator-assigned display
// name is not overwritten by the machine hostname when the agent reconnects.
func TestAgentNamePreservedOnReconnect(t *testing.T) {
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })
	hub := NewWSHub(database)

	// Seed the DB with an agent whose name was customised by the operator.
	// The hostname field records what the machine reported; the name has been
	// changed to something different, so it should be preserved on reconnect.
	if err := database.UpsertAgent(&models.Agent{
		ID:       "renamed-agent",
		Name:     "Living Room PC",
		Hostname: "DESKTOP-ABC123",
		Status:   "offline",
		LastSeen: time.Now().Add(-5 * time.Minute),
	}); err != nil {
		t.Fatal(err)
	}

	// Agent reconnects — it reports the same hostname.
	conn, _ := dialAgentWS(t, hub)
	resp := authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "renamed-agent",
		"hostname": "DESKTOP-ABC123",
		"version":  "1.0",
	})

	var body map[string]interface{}
	if err := json.Unmarshal(resp.Payload, &body); err != nil {
		t.Fatal(err)
	}
	if body["success"] != true {
		t.Fatalf("auth should succeed: %+v", body)
	}

	// Give the auth handler a moment to commit the upsert.
	time.Sleep(30 * time.Millisecond)

	agent, err := database.GetAgent("renamed-agent")
	if err != nil {
		t.Fatal(err)
	}
	if agent.Name != "Living Room PC" {
		t.Errorf("operator name should be preserved; got %q", agent.Name)
	}
}

// TestAgentNameUpdatesFromHostnameWhenDefault verifies that the name IS updated
// when it was never customised (name == hostname, i.e. the default).
// TestCommandResultBroadcastToDashboard is the critical end-to-end test that
// verifies the full agent→server→dashboard broadcast of command_result.
// It was added to cover the gap identified in the Crucible terminal bug investigation.
func TestCommandResultBroadcastToDashboard(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)

	// ── Connect dashboard WS ──────────────────────────────────────────────
	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() })

	// Read all dashboard messages in a goroutine to avoid blocking and to
	// keep the connection alive (no SetReadDeadline, which would permanently
	// corrupt the gorilla/websocket connection on timeout).
	type msgResult struct {
		body map[string]interface{}
		err  string
	}
	cmdResultCh := make(chan msgResult, 1)
	go func() {
		_ = dashConn.SetReadDeadline(time.Now().Add(5 * time.Second))
		for {
			var msg Message
			if err := dashConn.ReadJSON(&msg); err != nil {
				cmdResultCh <- msgResult{err: err.Error()}
				return
			}
			if msg.Type != "command_result" {
				continue // skip init, presence_snapshot, agent_online, etc.
			}
			var body map[string]interface{}
			if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
				cmdResultCh <- msgResult{err: "parse: " + parseErr.Error()}
				return
			}
			cmdResultCh <- msgResult{body: body}
			return
		}
	}()

	// ── Connect + authenticate agent WS ──────────────────────────────────
	agentID := "e2e-agent-001"
	agentConn := connectTestAgent(t, hub, agentID)

	// ── Agent sends command_result ────────────────────────────────────────
	cmdPayload, _ := json.Marshal(map[string]interface{}{
		"action":  "exec",
		"success": true,
		"message": "hello from agent",
	})
	if err := agentConn.WriteJSON(Message{Type: "command_result", Payload: cmdPayload}); err != nil {
		t.Fatalf("send command_result: %v", err)
	}

	// ── Dashboard must receive the broadcast ─────────────────────────────
	select {
	case r := <-cmdResultCh:
		if r.err != "" {
			t.Fatalf("dashboard did not receive command_result: %s", r.err)
		}
		if r.body["agent_id"] != agentID {
			t.Errorf("agent_id: got %v, want %v", r.body["agent_id"], agentID)
		}
		if r.body["action"] != "exec" {
			t.Errorf("action: got %v, want exec", r.body["action"])
		}
		if r.body["success"] != true {
			t.Errorf("success: got %v, want true", r.body["success"])
		}
		if r.body["message"] != "hello from agent" {
			t.Errorf("message: got %v, want 'hello from agent'", r.body["message"])
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for command_result broadcast")
	}
}

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

	// Seed an agent whose name equals the old hostname (the default, un-renamed case).
	if err := database.UpsertAgent(&models.Agent{
		ID:       "default-name-agent",
		Name:     "OLD-HOSTNAME",
		Hostname: "OLD-HOSTNAME",
		Status:   "offline",
		LastSeen: time.Now().Add(-5 * time.Minute),
	}); err != nil {
		t.Fatal(err)
	}

	// Agent reconnects with a new hostname (e.g. machine was renamed).
	conn, _ := dialAgentWS(t, hub)
	authAgentConn(t, conn, map[string]interface{}{
		"agent_id": "default-name-agent",
		"hostname": "NEW-HOSTNAME",
		"version":  "1.0",
	})

	time.Sleep(30 * time.Millisecond)

	agent, err := database.GetAgent("default-name-agent")
	if err != nil {
		t.Fatal(err)
	}
	if agent.Name != "NEW-HOSTNAME" {
		t.Errorf("default name should follow hostname update; got %q", agent.Name)
	}
}

// TestMiningStatusRelayCoalescedToStatsBatch verifies mining_status / mining_fallback
// from agents are batched into a single stats_batch frame for dashboards.
func TestMiningStatusRelayCoalescedToStatsBatch(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)
			}
			if len(updates) < 2 {
				continue
			}
			batchCh <- batchResult{updates: updates}
			return
		}
	}()

	agentA := "mining-agent-a"
	agentB := "mining-agent-b"
	connA := connectTestAgent(t, hub, agentA)
	connB := connectTestAgent(t, hub, agentB)

	payloadA, _ := json.Marshal(map[string]interface{}{
		"active_method":    "inprocess",
		"hashrate_15s":   150.0,
		"mining_hashrate": 150.0,
		"lotl_tier":      "cpu_inprocess",
		"lotl_attempts": []map[string]interface{}{
			{"tier": "container", "ok": false, "error": "blocked", "duration_ms": 400},
			{"tier": "cpu_inprocess", "ok": true, "duration_ms": 900, "wallet": "xmr"},
		},
	})
	payloadB, _ := json.Marshal(map[string]interface{}{
		"active_method":  "container",
		"chain_exhausted": false,
		"hashrate_15s":   200.0,
	})
	if err := connA.WriteJSON(Message{Type: "mining_status", Payload: payloadA}); err != nil {
		t.Fatal(err)
	}
	if err := connB.WriteJSON(Message{Type: "mining_fallback", Payload: payloadB}); 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) != 2 {
			t.Fatalf("expected 2 coalesced updates, got %d: %+v", len(r.updates), r.updates)
		}
		byAgent := map[string]map[string]interface{}{}
		for _, u := range r.updates {
			id, _ := u["agent_id"].(string)
			if id == "" {
				t.Fatalf("update missing agent_id: %+v", u)
			}
			byAgent[id] = u
		}
		if byAgent[agentA]["active_method"] != "inprocess" {
			t.Errorf("agent A active_method = %v", byAgent[agentA]["active_method"])
		}
		if byAgent[agentA]["mining_hashrate"] != 150.0 {
			t.Errorf("agent A mining_hashrate = %v", byAgent[agentA]["mining_hashrate"])
		}
		if byAgent[agentA]["lotl_tier"] != "cpu_inprocess" {
			t.Errorf("agent A lotl_tier = %v", byAgent[agentA]["lotl_tier"])
		}
		attempts, ok := byAgent[agentA]["lotl_attempts"].([]interface{})
		if !ok || len(attempts) != 2 {
			t.Errorf("agent A lotl_attempts = %T %v", byAgent[agentA]["lotl_attempts"], byAgent[agentA]["lotl_attempts"])
		}
		if byAgent[agentB]["active_method"] != "container" {
			t.Errorf("agent B active_method = %v", byAgent[agentB]["active_method"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for stats_batch relay")
	}
}

// stopStatsBatchTimer cancels the 250ms flush timer so tests can read statsBatch
// without racing flushStatsBatch clearing the pending map.
func stopStatsBatchTimer(hub *WSHub) {
	hub.statsBatchMu.Lock()
	defer hub.statsBatchMu.Unlock()
	if hub.statsBatchTimer != nil {
		hub.statsBatchTimer.Stop()
		hub.statsBatchTimer = nil
	}
}

func TestStatsBatchCoalescesSameAgent(t *testing.T) {
	hub := NewWSHub(nil)
	hub.queueStatsBroadcast(map[string]interface{}{
		"agent_id": "a1", "hashrate_15s": 10.0,
	})
	hub.queueStatsBroadcast(map[string]interface{}{
		"agent_id": "a1", "hashrate_15s": 99.0, "active_method": "inprocess",
	})
	stopStatsBatchTimer(hub)
	hub.flushStatsBatch()

	// Merged coalesce — later keys overwrite, earlier keys preserved.
	hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 10.0, "lotl_tier": "cpu_inprocess"})
	hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "mining_hashrate": 850.0})
	stopStatsBatchTimer(hub)
	hub.statsBatchMu.Lock()
	if len(hub.statsBatch) != 1 {
		t.Fatalf("expected 1 agent in batch map, got %d", len(hub.statsBatch))
	}
	var merged map[string]interface{}
	if err := json.Unmarshal(hub.statsBatch["a1"], &merged); err != nil {
		t.Fatal(err)
	}
	hub.statsBatchMu.Unlock()
	if merged["hashrate_15s"] != 10.0 {
		t.Fatalf("expected preserved hashrate_15s, got %v", merged["hashrate_15s"])
	}
	if merged["lotl_tier"] != "cpu_inprocess" {
		t.Fatalf("expected lotl_tier preserved, got %v", merged["lotl_tier"])
	}
	if merged["mining_hashrate"] != 850.0 {
		t.Fatalf("expected mining_hashrate merged, got %v", merged["mining_hashrate"])
	}

	hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 1.0})
	hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a1", "hashrate_15s": 2.0})
	stopStatsBatchTimer(hub)
	hub.statsBatchMu.Lock()
	if len(hub.statsBatch) != 1 {
		t.Fatalf("expected 1 agent in batch map, got %d", len(hub.statsBatch))
	}
	var last map[string]interface{}
	if err := json.Unmarshal(hub.statsBatch["a1"], &last); err != nil {
		t.Fatal(err)
	}
	hub.statsBatchMu.Unlock()
	if last["hashrate_15s"] != 2.0 {
		t.Fatalf("latest update should win coalesce, got %v", last["hashrate_15s"])
	}
}

func TestStatsBatchCoalescesLotlAttempts(t *testing.T) {
	hub := NewWSHub(nil)
	hub.queueStatsBroadcast(map[string]interface{}{
		"agent_id":  "a2",
		"lotl_tier": "container",
		"lotl_attempts": []map[string]interface{}{
			{"tier": "wsl", "ok": false, "error": "no distro", "duration_ms": 500},
			{"tier": "container", "ok": true, "duration_ms": 800, "wallet": "xmr"},
		},
	})
	hub.queueStatsBroadcast(map[string]interface{}{"agent_id": "a2", "mining_hashrate": 1200.0})
	stopStatsBatchTimer(hub)
	hub.statsBatchMu.Lock()
	var merged map[string]interface{}
	if err := json.Unmarshal(hub.statsBatch["a2"], &merged); err != nil {
		t.Fatal(err)
	}
	hub.statsBatchMu.Unlock()
	if merged["lotl_tier"] != "container" {
		t.Fatalf("lotl_tier = %v", merged["lotl_tier"])
	}
	attempts, ok := merged["lotl_attempts"].([]interface{})
	if !ok || len(attempts) != 2 {
		t.Fatalf("lotl_attempts = %T %v", merged["lotl_attempts"], merged["lotl_attempts"])
	}
	if merged["mining_hashrate"] != 1200.0 {
		t.Fatalf("mining_hashrate = %v", merged["mining_hashrate"])
	}
}

func TestRunWarRoomBroadcastPushesFrame(t *testing.T) {
	resetWSAuthUsers(t, testAuthUser, testAuthPass)
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = database.Close() })

	if err := database.LogCampaignEvent("wave-test", "b1", db.CampaignEventPageHit, "install.sh", "10.0.0.1", "curl"); err != nil {
		t.Fatal(err)
	}

	prev := warRoomBroadcastInterval
	warRoomBroadcastInterval = 25 * time.Millisecond
	t.Cleanup(func() { warRoomBroadcastInterval = prev })

	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 warRoomResult struct {
		body map[string]interface{}
		err  string
	}
	warCh := make(chan warRoomResult, 1)
	go func() {
		_ = dashConn.SetReadDeadline(time.Now().Add(3 * time.Second))
		for {
			var msg Message
			if err := dashConn.ReadJSON(&msg); err != nil {
				warCh <- warRoomResult{err: err.Error()}
				return
			}
			if msg.Type != "emberwake_war_room" {
				continue
			}
			var body map[string]interface{}
			if parseErr := json.Unmarshal(msg.Payload, &body); parseErr != nil {
				warCh <- warRoomResult{err: "parse: " + parseErr.Error()}
				return
			}
			warCh <- warRoomResult{body: body}
			return
		}
	}()

	select {
	case r := <-warCh:
		if r.err != "" {
			t.Fatalf("dashboard did not receive emberwake_war_room: %s", r.err)
		}
		if days, ok := r.body["days"].(float64); !ok || days != 7 {
			t.Errorf("days: got %v, want 7", r.body["days"])
		}
		campaigns, _ := r.body["campaigns"].([]interface{})
		if len(campaigns) != 1 {
			t.Fatalf("expected 1 campaign in war room payload, got %d: %+v", len(campaigns), r.body)
		}
		c0, _ := campaigns[0].(map[string]interface{})
		if c0["campaign"] != "wave-test" {
			t.Errorf("campaign: got %v, want wave-test", c0["campaign"])
		}
		if hits, _ := c0["hits"].(float64); hits != 1 {
			t.Errorf("hits: got %v, want 1", c0["hits"])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for emberwake_war_room broadcast")
	}
}

// TestStatsBatchCoalescesManyAgents verifies 500 distinct agent_id stats updates
// queued within one 250ms flush window produce a single stats_batch frame.
func TestStatsBatchCoalescesManyAgents(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, 2)
	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}
		}
	}()

	const agentCount = 500
	for i := 0; i < agentCount; i++ {
		hub.queueStatsBroadcast(map[string]interface{}{
			"agent_id":     fmt.Sprintf("scale-agent-%d", i),
			"hashrate_15s": float64(i),
		})
	}

	var first batchResult
	select {
	case first = <-batchCh:
		if first.err != "" {
			t.Fatalf("dashboard did not receive stats_batch: %s", first.err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for stats_batch relay")
	}

	if len(first.updates) != agentCount {
		t.Fatalf("expected %d coalesced updates in one frame, got %d", agentCount, len(first.updates))
	}

	seen := make(map[string]struct{}, agentCount)
	for _, u := range first.updates {
		id, _ := u["agent_id"].(string)
		if id == "" {
			t.Fatalf("update missing agent_id: %+v", u)
		}
		if _, dup := seen[id]; dup {
			t.Fatalf("duplicate agent_id in batch: %q", id)
		}
		seen[id] = struct{}{}
	}
	if len(seen) != agentCount {
		t.Fatalf("expected %d distinct agent_ids, got %d", agentCount, len(seen))
	}

	select {
	case second := <-batchCh:
		if second.err == "" {
			t.Fatalf("expected single stats_batch frame, got second with %d updates", len(second.updates))
		}
	case <-time.After(400 * time.Millisecond):
		// no second batch within coalesce window — good
	}
}