frontend/components/RelayAPI.js

// RelayAPI.js - Next.js Wrapper for the WebSocket Server

import { useEffect, useState, useCallback, useRef } from "react";

const useWebSocket = (userId) => {
  const [socket, setSocket] = useState(null);
  const [messages, setMessages] = useState([]);
  const [connectedChannels, setConnectedChannels] = useState(new Set());
  const reconnectInterval = useRef(null);

  const connect = useCallback(() => {
    const ws = new WebSocket(`${process.env.NEXT_PUBLIC_RELAY_URL}/ws?user_id=${userId}`);

    ws.onopen = () => {
      console.debug("WebSocket connected.");
      clearInterval(reconnectInterval.current);
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        console.debug("Received message:", data);
        setMessages((prev) => [...prev, data]);
        console.debug(messages);
      } catch (error) {
        console.debug("Failed to parse WebSocket message:", event.data);
      }
    };

    ws.onclose = () => {
      console.warn("WebSocket closed. Reconnecting...");
      reconnectInterval.current = setInterval(() => connect(), 5000);
    };

    ws.onerror = (error) => {
      console.debug("WebSocket error:", error);
      ws.close();
    };

    setSocket(ws);
  }, [userId]);

  useEffect(() => {
    if (!userId) return;

    connect();

    return () => {
      clearInterval(reconnectInterval.current);
      setConnectedChannels(new Set());
      if (socket) socket.close();
    };
  }, [connect, userId]);

  const sendMessage = useCallback(
    (action, channel, message, targetUserId = null) => {
      if (!socket || socket.readyState !== WebSocket.OPEN) {
        console.warn("WebSocket is not connected.");
        return;
      }

      const payload = { action, channel, message, sender: userId };
      if (targetUserId) payload.target_user_id = targetUserId;

      try {
        socket.send(JSON.stringify(payload));
        console.debug("Sent message:", payload);
      } catch (error) {
        console.debug("Failed to send WebSocket message:", error);
      }
    },
    [socket, userId]
  );

  const joinChannel = useCallback(
    (channel) => {
      if (connectedChannels.has(channel)) return;

      sendMessage("join", channel);
      setConnectedChannels((prev) => new Set([...prev, channel]));
    },
    [sendMessage, connectedChannels]
  );

  const leaveChannel = useCallback(
    (channel) => {
      if (!connectedChannels.has(channel)) return;

      sendMessage("leave", channel);
      setConnectedChannels((prev) => {
        const updated = new Set(prev);
        updated.delete(channel);
        return updated;
      });
    },
    [sendMessage, connectedChannels]
  );

  const broadcast = useCallback(
    (channel, message) => {
      sendMessage("broadcast", channel, message);
    },
    [sendMessage]
  );

  const sendToUser = useCallback(
    (targetUserId, message) => {
      sendMessage("send_to_user", null, message, targetUserId);
    },
    [sendMessage]
  );

  return { messages, joinChannel, leaveChannel, broadcast, sendToUser };
};

export const RelayAPIProvider = ({ userId=null, children }) => {
  if (userId === null) {
    userId = localStorage.getItem("u_id");
    console.log("inside context")
  }
  const contextValue = useWebSocket(userId);

  return (
    <RelayAPIContext.Provider value={contextValue}>{children}</RelayAPIContext.Provider>
  );
};

import React, { useContext } from "react";

const RelayAPIContext = React.createContext(null);

export const useRelayAPI = () => {
  const context = useContext(RelayAPIContext);
  if (!context) {
    throw new Error("useRelayAPI must be used within a RelayAPIProvider");
  }
  return context;
};