pulsenode/lib/ws/client.ts

"use client";

import { useEffect, useState } from "react";
import type { Config } from "@/lib/config/schema";
import type { WidgetResult } from "@/lib/types/widget-result";

interface Envelope {
  topic: string;
  type: string;
  ts: number;
  data: unknown;
}

type TopicListener = (envelope: Envelope) => void;
export type ConnectionStatus = "connecting" | "connected" | "reconnecting";
type ConnectionListener = (status: ConnectionStatus) => void;

let sharedSocket: WebSocket | null = null;
let refCount = 0;
let reconnectAttempts = 0;
let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
let connectionStatus: ConnectionStatus = "connecting";
const topicListeners = new Map<string, Set<TopicListener>>();
const connectionListeners = new Set<ConnectionListener>();
const pendingWidgetSubscriptions = new Set<string>();

function setConnectionStatus(status: ConnectionStatus): void {
  connectionStatus = status;
  for (const listener of connectionListeners) listener(status);
}

function getSocketUrl(): string {
  const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
  return `${protocol}//${window.location.host}/ws`;
}

function scheduleReconnect(): void {
  if (reconnectTimer || refCount === 0) return;
  setConnectionStatus("reconnecting");
  const delay = Math.min(30_000, 1_000 * 2 ** reconnectAttempts);
  reconnectAttempts += 1;
  reconnectTimer = setTimeout(() => {
    reconnectTimer = null;
    if (refCount > 0) ensureSocket();
  }, delay);
}

function ensureSocket(): WebSocket {
  if (sharedSocket && sharedSocket.readyState <= WebSocket.OPEN) {
    return sharedSocket;
  }

  const socket = new WebSocket(getSocketUrl());
  sharedSocket = socket;

  socket.addEventListener("open", () => {
    reconnectAttempts = 0;
    setConnectionStatus("connected");
    for (const widgetId of pendingWidgetSubscriptions) {
      socket.send(JSON.stringify({ action: "subscribe", widgetId }));
    }
  });

  socket.addEventListener("message", (event) => {
    let envelope: Envelope;
    try {
      envelope = JSON.parse(event.data);
    } catch {
      return;
    }
    const listeners = topicListeners.get(envelope.topic);
    if (!listeners) return;
    for (const listener of listeners) listener(envelope);
  });

  socket.addEventListener("close", () => {
    if (sharedSocket === socket) sharedSocket = null;
    scheduleReconnect();
  });

  return socket;
}

function addTopicListener(topic: string, listener: TopicListener): () => void {
  let listeners = topicListeners.get(topic);
  if (!listeners) {
    listeners = new Set();
    topicListeners.set(topic, listeners);
  }
  listeners.add(listener);
  return () => {
    listeners?.delete(listener);
    if (listeners && listeners.size === 0) topicListeners.delete(topic);
  };
}

export function useWidgetSubscription(widgetId: string): WidgetResult | null {
  const [result, setResult] = useState<WidgetResult | null>(null);

  useEffect(() => {
    refCount += 1;
    pendingWidgetSubscriptions.add(widgetId);
    const socket = ensureSocket();
    const topic = `widget:${widgetId}`;

    const removeListener = addTopicListener(topic, (envelope) => {
      if (envelope.type === "result") setResult(envelope.data as WidgetResult);
    });

    if (socket.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify({ action: "subscribe", widgetId }));
    }

    return () => {
      refCount -= 1;
      removeListener();
      if (!topicListeners.has(topic)) {
        pendingWidgetSubscriptions.delete(widgetId);
        if (socket.readyState === WebSocket.OPEN) {
          socket.send(JSON.stringify({ action: "unsubscribe", widgetId }));
        }
      }
    };
  }, [widgetId]);

  return result;
}

export function useConfigSubscription(initial: Config): Config {
  const [config, setConfig] = useState(initial);

  useEffect(() => {
    refCount += 1;
    ensureSocket();

    const removeListener = addTopicListener("config", (envelope) => {
      if (envelope.type === "config:update") setConfig(envelope.data as Config);
    });

    return () => {
      refCount -= 1;
      removeListener();
    };
  }, []);

  return config;
}

export function useConnectionStatus(): ConnectionStatus {
  const [status, setStatus] = useState<ConnectionStatus>(connectionStatus);

  useEffect(() => {
    refCount += 1;
    // ensureSocket() can synchronously open a connection another component
    // already established, so the module-level status may have moved past
    // what this hook's initial useState captured - reconcile once before
    // subscribing to future changes.
    ensureSocket();
    connectionListeners.add(setStatus);
    // eslint-disable-next-line react-hooks/set-state-in-effect
    setStatus(connectionStatus);

    return () => {
      refCount -= 1;
      connectionListeners.delete(setStatus);
    };
  }, []);

  return status;
}

export interface ConfigToast {
  id: number;
  tone: "ok" | "error";
  message: string;
}

let toastCounter = 0;

export function useConfigToasts(): ConfigToast[] {
  const [toasts, setToasts] = useState<ConfigToast[]>([]);

  useEffect(() => {
    refCount += 1;
    ensureSocket();

    const removeListener = addTopicListener("config", (envelope) => {
      let toast: Omit<ConfigToast, "id"> | null = null;
      if (envelope.type === "config:update") {
        toast = { tone: "ok", message: "config.yml reloaded" };
      } else if (envelope.type === "config:error") {
        toast = { tone: "error", message: "config.yml has an error — showing last valid config" };
      }
      if (!toast) return;

      toastCounter += 1;
      const id = toastCounter;
      setToasts((prev) => [...prev, { ...toast!, id }]);
      setTimeout(() => setToasts((prev) => prev.filter((t) => t.id !== id)), 4000);
    });

    return () => {
      refCount -= 1;
      removeListener();
    };
  }, []);

  return toasts;
}