AetherForge/server/web/src/hooks/useWebSocket.ts

import { useEffect, useRef, useCallback, useState } from 'react';
import type { WSMessage, Agent, Share, FleetAlert, PoolStatus, AIActivityEntry } from '../types';

interface DashboardInit {
  agents: Agent[];
}

interface UseWebSocketReturn {
  isConnected: boolean;
  agents: Agent[];
  recentShares: Share[];
  fleetAlerts: FleetAlert[];
  poolStatus: PoolStatus[];
  aiActivity: AIActivityEntry[];
  agentLogs: Record<string, string>;
}

export function useWebSocket(): UseWebSocketReturn {
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const unmounted = useRef(false);
  const [isConnected, setIsConnected] = useState(false);
  const [agents, setAgents] = useState<Agent[]>([]);
  const [recentShares, setRecentShares] = useState<Share[]>([]);
  const [fleetAlerts, setFleetAlerts] = useState<FleetAlert[]>([]);
  const [poolStatus, setPoolStatus] = useState<PoolStatus[]>([]);
  const [aiActivity, setAiActivity] = useState<AIActivityEntry[]>([]);
  const [agentLogs, setAgentLogs] = useState<Record<string, string>>({});

  const connect = useCallback(() => {
    if (unmounted.current) return;

    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const wsUrl = `${protocol}//${window.location.host}/ws/dashboard`;

    const ws = new WebSocket(wsUrl);
    wsRef.current = ws;

    ws.onopen = () => {
      if (!unmounted.current) setIsConnected(true);
    };

    ws.onclose = () => {
      if (unmounted.current) return;
      setIsConnected(false);
      reconnectTimer.current = setTimeout(connect, 3000);
    };

    ws.onerror = () => {
      ws.close();
    };

    ws.onmessage = (event) => {
      try {
        const msg: WSMessage = JSON.parse(event.data);

        switch (msg.type) {
          case 'init': {
            const data = msg.payload as DashboardInit;
            if (data.agents) setAgents(data.agents);
            break;
          }
          case 'agent_online': {
            const agent = msg.payload as Agent;
            setAgents((prev) => {
              const idx = prev.findIndex((a) => a.id === agent.id);
              if (idx >= 0) {
                const updated = [...prev];
                updated[idx] = agent;
                return updated;
              }
              return [...prev, agent];
            });
            break;
          }
          case 'agent_offline': {
            const { agent_id } = msg.payload as { agent_id: string };
            setAgents((prev) =>
              prev.map((a) =>
                a.id === agent_id ? { ...a, status: 'offline' as const } : a
              )
            );
            break;
          }
          case 'stats_update': {
            const update = msg.payload as {
              agent_id: string;
              hashrate_15s: number;
              hashrate_1m: number;
              hashrate_15m: number;
              cpu_usage_pct: number;
            };
            setAgents((prev) =>
              prev.map((a) =>
                a.id === update.agent_id
                  ? {
                      ...a,
                      hashrate_15s: update.hashrate_15s,
                      hashrate_1m: update.hashrate_1m,
                      hashrate_15m: update.hashrate_15m,
                      cpu_usage_pct: update.cpu_usage_pct,
                    }
                  : a
              )
            );
            break;
          }
          case 'new_share': {
            const share = msg.payload as Share;
            setRecentShares((prev) => [share, ...prev].slice(0, 50));
            break;
          }
          case 'fleet_alert': {
            const alert = msg.payload as FleetAlert;
            setFleetAlerts((prev) => [alert, ...prev].slice(0, 20));
            break;
          }
          case 'pool_status': {
            const pools = msg.payload as PoolStatus[];
            if (Array.isArray(pools)) setPoolStatus(pools);
            break;
          }
          case 'ai_activity': {
            const entry = msg.payload as AIActivityEntry;
            setAiActivity((prev) => {
              const idx = prev.findIndex((a) => a.agent_id === entry.agent_id);
              if (idx >= 0) {
                const next = [...prev];
                next[idx] = entry;
                return next;
              }
              return [...prev, entry];
            });
            break;
          }
          case 'agent_log': {
            const { agent_id, content } = msg.payload as { agent_id: string; content: string };
            if (agent_id) {
              setAgentLogs((prev) => ({ ...prev, [agent_id]: content }));
            }
            break;
          }
        }
      } catch (err) {
        console.error('Failed to parse WebSocket message:', err);
      }
    };
  }, []);

  useEffect(() => {
    unmounted.current = false;
    connect();
    return () => {
      unmounted.current = true;
      if (reconnectTimer.current) {
        clearTimeout(reconnectTimer.current);
      }
      const ws = wsRef.current;
      if (ws) {
        ws.onclose = null;
        ws.close();
      }
    };
  }, [connect]);

  return { isConnected, agents, recentShares, fleetAlerts, poolStatus, aiActivity, agentLogs };
}