saltylab-firmware/ui/social-bot/src/hooks/useFleet.js

/**
 * useFleet.js — Multi-robot fleet connection manager.
 *
 * Maintains one ROSLIB.Ros WebSocket connection per robot.
 * Aggregates fleet-summary telemetry (battery, pose, alerts, mode) for each.
 * Persists robot registry to localStorage.
 *
 * API:
 *   const { robots, addRobot, removeRobot, updateRobot,
 *           connections, robotData,
 *           subscribeRobot, publishRobot, scanSubnet } = useFleet();
 *
 * Robot shape:
 *   { id, name, variant, host, port, videoPort, videoTopic }
 *   variant: 'balance' | 'rover' | 'tank' | 'social'
 *   videoTopic: '/camera/panoramic/compressed' (default)
 */

import { useState, useEffect, useRef, useCallback } from 'react';
import ROSLIB from 'roslib';

// ── Storage ────────────────────────────────────────────────────────────────────

const STORAGE_KEY = 'saltybot_fleet_v2';

const DEFAULT_ROBOTS = [
  {
    id: 'salty-01',
    name: 'SaltyBot #1',
    variant: 'balance',
    host: '192.168.1.100',
    port: 9090,
    videoPort: 8080,
    videoTopic: '/camera/panoramic/compressed',
  },
];

function loadRobots() {
  try {
    const saved = JSON.parse(localStorage.getItem(STORAGE_KEY));
    return Array.isArray(saved) && saved.length > 0 ? saved : DEFAULT_ROBOTS;
  } catch {
    return DEFAULT_ROBOTS;
  }
}

// ── Robot data structure ───────────────────────────────────────────────────────

const EMPTY_DATA = () => ({
  state:       null,   // 'DISARMED'|'ARMED'|'TILT FAULT'
  stm32Mode:   null,   // 'MANUAL'|'ASSISTED'
  controlMode: null,   // 'RC'|'RAMP_TO_AUTO'|'AUTO'|'RAMP_TO_RC'
  blendAlpha:  0,
  pipeline:    null,   // 'idle'|'listening'|'thinking'|'speaking'
  pitch:       0,
  motorCmd:    0,
  battery:     { voltage: 0, soc: 0 },
  pose:        null,   // { x, y, yaw }
  alerts:      [],     // [{ level, name, message }]
  lastSeen:    null,
});

// ── Hook ──────────────────────────────────────────────────────────────────────

export function useFleet() {
  const [robots, setRobots]           = useState(loadRobots);
  const [connections, setConnections] = useState({});
  const [robotData, setRobotData]     = useState(() =>
    Object.fromEntries(loadRobots().map(r => [r.id, EMPTY_DATA()]))
  );
  const [scanning, setScanning]       = useState(false);

  const rosRef  = useRef({});  // { [id]: ROSLIB.Ros }
  const subsRef = useRef({});  // { [id]: ROSLIB.Topic[] }

  // Persist robots
  useEffect(() => {
    localStorage.setItem(STORAGE_KEY, JSON.stringify(robots));
  }, [robots]);

  // ── Connection lifecycle ────────────────────────────────────────────────────

  useEffect(() => {
    const currentIds = new Set(robots.map(r => r.id));

    // Disconnect removed robots
    Object.keys(rosRef.current).forEach(id => {
      if (!currentIds.has(id)) {
        subsRef.current[id]?.forEach(t => { try { t.unsubscribe(); } catch {} });
        delete subsRef.current[id];
        try { rosRef.current[id].close(); } catch {}
        delete rosRef.current[id];
        setConnections(c => { const n = { ...c }; delete n[id]; return n; });
      }
    });

    // Connect new robots
    robots.forEach(robot => {
      if (rosRef.current[robot.id]) return;

      const ros = new ROSLIB.Ros({ url: `ws://${robot.host}:${robot.port}` });
      rosRef.current[robot.id] = ros;
      subsRef.current[robot.id] = [];

      ros.on('connection', () => {
        setConnections(c => ({ ...c, [robot.id]: { connected: true, error: null } }));
        _setupSubs(robot.id, ros);
      });
      ros.on('error', err => {
        setConnections(c => ({
          ...c,
          [robot.id]: { connected: false, error: String(err?.message ?? err ?? 'error') }
        }));
      });
      ros.on('close', () => {
        setConnections(c => ({
          ...c,
          [robot.id]: { connected: false, error: null }
        }));
      });
    });

    return () => {
      Object.values(subsRef.current).forEach(subs =>
        subs.forEach(t => { try { t.unsubscribe(); } catch {} })
      );
      Object.values(rosRef.current).forEach(ros => { try { ros.close(); } catch {} });
      rosRef.current = {};
      subsRef.current = {};
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [robots]);

  // ── Fleet-summary subscriptions ────────────────────────────────────────────

  function _sub(id, ros, name, type, cb) {
    const t = new ROSLIB.Topic({ ros, name, messageType: type });
    t.subscribe(cb);
    subsRef.current[id] = [...(subsRef.current[id] ?? []), t];
  }

  function _update(id, patch) {
    setRobotData(prev => ({
      ...prev,
      [id]: { ...(prev[id] ?? EMPTY_DATA()), ...patch },
    }));
  }

  function _setupSubs(id, ros) {
    // Balance state
    _sub(id, ros, '/saltybot/balance_state', 'std_msgs/String', msg => {
      try {
        const d = JSON.parse(msg.data);
        _update(id, {
          state:     d.state,
          stm32Mode: d.mode,
          pitch:     d.pitch_deg ?? 0,
          motorCmd:  d.motor_cmd ?? 0,
          lastSeen:  Date.now(),
        });
      } catch {}
    });

    // Control mode (RC ↔ AUTO blend)
    _sub(id, ros, '/saltybot/control_mode', 'std_msgs/String', msg => {
      try {
        const d = JSON.parse(msg.data);
        _update(id, { controlMode: d.mode, blendAlpha: d.blend_alpha ?? 0 });
      } catch {}
    });

    // Odometry → pose
    _sub(id, ros, '/odom', 'nav_msgs/Odometry', msg => {
      const p = msg.pose?.pose;
      if (!p) return;
      const o = p.orientation;
      const yaw = Math.atan2(
        2 * (o.w * o.z + o.x * o.y),
        1 - 2 * (o.y * o.y + o.z * o.z)
      );
      _update(id, { pose: { x: p.position.x, y: p.position.y, yaw }, lastSeen: Date.now() });
    });

    // Diagnostics → battery + alerts
    _sub(id, ros, '/diagnostics', 'diagnostic_msgs/DiagnosticArray', msg => {
      const alerts = [];
      const kv = {};
      for (const status of msg.status ?? []) {
        if (status.level > 0) {
          alerts.push({ level: status.level, name: status.name, message: status.message });
        }
        for (const pair of status.values ?? []) kv[pair.key] = pair.value;
      }
      const patch = { alerts };
      if (kv.battery_voltage_v != null) {
        const v = parseFloat(kv.battery_voltage_v);
        patch.battery = {
          voltage: v,
          soc: Math.max(0, Math.min(100, ((v - 12) / (16.8 - 12)) * 100)),
        };
      }
      _update(id, patch);
    });

    // Social pipeline state (optional)
    _sub(id, ros, '/social/orchestrator/state', 'std_msgs/String', msg => {
      try {
        _update(id, { pipeline: JSON.parse(msg.data)?.state ?? null });
      } catch {}
    });
  }

  // ── Public API ─────────────────────────────────────────────────────────────

  const addRobot = useCallback(robot => {
    const id = robot.id ?? `bot-${Date.now()}`;
    const full = {
      videoPort: 8080,
      videoTopic: '/camera/panoramic/compressed',
      ...robot,
      id,
    };
    setRobots(r => [...r.filter(b => b.id !== id), full]);
    setRobotData(d => ({ ...d, [id]: EMPTY_DATA() }));
  }, []);

  const removeRobot = useCallback(id => {
    setRobots(r => r.filter(b => b.id !== id));
    setRobotData(d => { const n = { ...d }; delete n[id]; return n; });
  }, []);

  const updateRobot = useCallback((id, patch) => {
    // Force reconnect by dropping the existing connection
    subsRef.current[id]?.forEach(t => { try { t.unsubscribe(); } catch {} });
    subsRef.current[id] = [];
    try { rosRef.current[id]?.close(); } catch {}
    delete rosRef.current[id];
    setRobots(r => r.map(b => b.id === id ? { ...b, ...patch } : b));
  }, []);

  /** Subscribe to a topic on a specific robot. Returns unsubscribe fn. */
  const subscribeRobot = useCallback((id, name, type, cb) => {
    const ros = rosRef.current[id];
    if (!ros) return () => {};
    const t = new ROSLIB.Topic({ ros, name, messageType: type });
    t.subscribe(cb);
    subsRef.current[id] = [...(subsRef.current[id] ?? []), t];
    return () => {
      try { t.unsubscribe(); } catch {}
      subsRef.current[id] = subsRef.current[id]?.filter(x => x !== t) ?? [];
    };
  }, []);

  /** Publish to a topic on a specific robot. */
  const publishRobot = useCallback((id, name, type, data) => {
    const ros = rosRef.current[id];
    if (!ros) return;
    const t = new ROSLIB.Topic({ ros, name, messageType: type });
    t.publish(new ROSLIB.Message(data));
  }, []);

  /** Send a single Nav2 goal pose to a robot. */
  const sendGoal = useCallback((id, x, y, yaw = 0) => {
    const cy2 = Math.cos(yaw / 2);
    const sy2 = Math.sin(yaw / 2);
    publishRobot(id, '/goal_pose', 'geometry_msgs/PoseStamped', {
      header: { frame_id: 'map', stamp: { sec: 0, nanosec: 0 } },
      pose: {
        position: { x, y, z: 0 },
        orientation: { x: 0, y: 0, z: sy2, w: cy2 },
      },
    });
  }, [publishRobot]);

  /** Send a patrol route (array of {x,y,yaw?}) to a robot via PoseArray → GPS follower. */
  const sendPatrol = useCallback((id, waypoints) => {
    publishRobot(id, '/outdoor/waypoints', 'geometry_msgs/PoseArray', {
      header: { frame_id: 'map', stamp: { sec: 0, nanosec: 0 } },
      poses: waypoints.map(({ x, y, yaw = 0 }) => {
        const cy2 = Math.cos(yaw / 2);
        const sy2 = Math.sin(yaw / 2);
        return {
          position: { x, y, z: 0 },
          orientation: { x: 0, y: 0, z: sy2, w: cy2 },
        };
      }),
    });
  }, [publishRobot]);

  /**
   * Probe a list of candidate URLs to find reachable rosbridge servers.
   * Reports each discovered robot back via the addRobot callback.
   */
  const scanSubnet = useCallback(async (baseIp, portRange = [9090, 9097]) => {
    setScanning(true);
    const [portMin, portMax] = portRange;
    const ipParts = baseIp.split('.');
    const prefix = ipParts.slice(0, 3).join('.');

    const results = [];
    const probes = [];
    for (let host4 = 1; host4 <= 254; host4++) {
      for (let port = portMin; port <= portMax; port++) {
        probes.push({ host: `${prefix}.${host4}`, port });
      }
    }

    // Probe in batches of 32 to avoid thundering-herd
    const BATCH = 32;
    for (let i = 0; i < probes.length; i += BATCH) {
      const batch = probes.slice(i, i + BATCH);
      await Promise.allSettled(
        batch.map(({ host, port }) =>
          new Promise(resolve => {
            const ws = new WebSocket(`ws://${host}:${port}`);
            const timer = setTimeout(() => { try { ws.close(); } catch {} resolve(null); }, 800);
            ws.onopen = () => {
              clearTimeout(timer);
              results.push({ host, port });
              ws.close();
              resolve({ host, port });
            };
            ws.onerror = () => { clearTimeout(timer); resolve(null); };
          })
        )
      );
    }

    setScanning(false);
    return results;
  }, []);

  return {
    robots,
    addRobot,
    removeRobot,
    updateRobot,
    connections,
    robotData,
    subscribeRobot,
    publishRobot,
    sendGoal,
    sendPatrol,
    scanning,
    scanSubnet,
  };
}