Copy

import { type HoneypotEvent } from "src/types";
import { URL } from "utils/utils";
import { onBeforeUnmount, ref } from "vue";

export type ConnectionStatus = "disconnected" | "connecting" | "open";

// Global state for shared WebSocket connection
const connectionStatus = ref<ConnectionStatus>("disconnected");
let socket: WebSocket | null = null;
let reconnectTimeout: ReturnType<typeof setTimeout> | null = null;
let manualDisconnect = false;
const listeners = new Set<(events: HoneypotEvent[]) => void>();

function buildWebSocketUrl(): string {
  const url = URL().replace("http", "ws");
  return `${url}/ws`;
}

function connectInternal() {
  if (
    socket?.readyState === WebSocket.OPEN ||
    connectionStatus.value === "connecting"
  ) {
    return;
  }

  connectionStatus.value = "connecting";
  const url = buildWebSocketUrl();
  const ws = new WebSocket(url);
  socket = ws;

  ws.onopen = () => {
    connectionStatus.value = "open";
    if (reconnectTimeout) {
      clearTimeout(reconnectTimeout);
      reconnectTimeout = null;
    }
  };

  ws.onclose = () => {
    connectionStatus.value = "disconnected";
    // Try to reconnect after a short delay.
    if (!manualDisconnect) {
      reconnectTimeout = setTimeout(() => connectInternal(), 2000);
    }
  };

  ws.onerror = () => {
    // Errors are followed by onclose; we just mark closed here.
    connectionStatus.value = "disconnected";
  };

  ws.onmessage = (ev: MessageEvent) => {
    try {
      const parsed = JSON.parse(ev.data);

      // Handle both single events and arrays of events
      const eventArray: HoneypotEvent[] = Array.isArray(parsed)
        ? (parsed as HoneypotEvent[])
        : [parsed as HoneypotEvent];

      // Filter fields for all events
      for (const data of eventArray) {
        delete data.fields?.body;
        delete data.fields?.port;
        delete data.fields?.remote_addr;
      }

      // Call all listeners with the entire batch
      listeners.forEach((listener) => listener(eventArray));
    } catch {
      // Ignore malformed messages.
    }
  };
}

function disconnectInternal() {
  manualDisconnect = true;
  if (reconnectTimeout) {
    clearTimeout(reconnectTimeout);
    reconnectTimeout = null;
  }
  if (socket) {
    socket.close();
    socket = null;
  }
}

export function useWebSocket(
  onMessage?: (events: HoneypotEvent[]) => void,
  maxEvents = 1000,
) {
  const events = ref<HoneypotEvent[]>([]);

  // Wrap the onMessage to also update local events ref
  const internalListener = (eventArray: HoneypotEvent[]) => {
    events.value.unshift(...eventArray);
    if (events.value.length > maxEvents) {
      events.value.length = maxEvents;
    }
    if (onMessage) {
      onMessage(eventArray);
    }
  };

  listeners.add(internalListener);

  onBeforeUnmount(() => {
    listeners.delete(internalListener);
    if (listeners.size === 0) {
      disconnectInternal();
    }
  });

  return {
    connectionStatus,
    events,
    connect: () => {
      manualDisconnect = false;
      connectInternal();
    },
    disconnect: disconnectInternal,
  };
}