Vitron-Front/app/hooks/useWebSocket.ts

import { useEffect, useState } from "react";

import { useAuthToken } from "./use-root-data";
import { getSocketBaseUrl } from "~/utils/api-config";

const useWebSocket = (threadId: string, userId: string | undefined) => {
  const [socket, setSocket] = useState<WebSocket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [newMessage, setNewMessage] = useState<NewMessageEvent | null>(null);
  const [readReceipt, setReadReceipt] = useState<ReadReceiptEvent | null>(null);
  const [messageDeleted, setMessageDeleted] =
    useState<MessageDeletedEvent | null>(null);

  const token = useAuthToken();
  const baseSocketPath = getSocketBaseUrl();
  // Initialize WebSocket connection
  useEffect(() => {
    if (!threadId || !userId || !token) return;

    // Use wss in production
    const wsUrl = `${baseSocketPath}/ws/chat/${threadId}/?token=${token}`;

    const ws = new WebSocket(wsUrl);

    ws.onopen = () => {
      console.log("WebSocket connection established");
      setIsConnected(true);
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data) as WebSocketMessage;

        if (data.type === "new_message") {
          setNewMessage(data as NewMessageEvent);
        } else if (data.type === "read_receipt") {
          setReadReceipt(data as ReadReceiptEvent);
        } else if (data.type === "message_deleted") {
          setMessageDeleted(data as MessageDeletedEvent);
        }
      } catch (error) {
        console.error("Error parsing WebSocket message:", error);
      }
    };

    ws.onclose = () => {
      console.log("WebSocket connection closed");
      setIsConnected(false);
    };

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

    setSocket(ws);

    // Clean up on unmount
    return () => {
      if (
        ws.readyState === WebSocket.OPEN ||
        ws.readyState === WebSocket.CONNECTING
      ) {
        ws.close();
      }
    };
  }, [threadId, userId, token]);

  // Send a message through WebSocket
  const sendWebSocketMessage = (
    content: string,
    replyToId?: string | null,
    replyContent?: string | null,
    imageUrl?: string | null,
  ) => {
    if (!socket || socket.readyState !== WebSocket.OPEN || !userId) {
      console.error("WebSocket is not connected");
      return false;
    }

    const message: SendMessagePayload = {
      type: "new_message",
      content,
      ...(replyToId && { reply_to_id: replyToId }),
      ...(replyContent && { reply_content: replyContent }),
      ...(imageUrl && { image_url: imageUrl }),
    };

    socket.send(JSON.stringify(message));
    return true;
  };

  // Mark messages as read
  const markMessagesAsRead = () => {
    if (!socket || socket.readyState !== WebSocket.OPEN || !userId) {
      console.error("WebSocket is not connected");
      return false;
    }

    const message: MarkReadPayload = {
      type: "mark_read",
      // reader_id: userId, // Remove this as it's not needed in client->server communication
    };

    socket.send(JSON.stringify(message));
    return true;
  };

  // Delete a message
  const deleteMessage = (messageId: string) => {
    if (!socket || socket.readyState !== WebSocket.OPEN) {
      console.error("WebSocket is not connected");
      return false;
    }

    const message: DeleteMessagePayload = {
      type: "delete_message",
      message_id: messageId,
    };

    socket.send(JSON.stringify(message));
    return true;
  };

  return {
    isConnected,
    sendWebSocketMessage,
    markMessagesAsRead,
    deleteMessage,
    newMessage,
    readReceipt,
    messageDeleted,
    // Reset message notification after handling
    resetNewMessage: () => setNewMessage(null),
    resetReadReceipt: () => setReadReceipt(null),
    resetMessageDeleted: () => setMessageDeleted(null),
  };
};

// WebSocket message types
interface WebSocketMessage {
  type: string;
  [key: string]: string | number | boolean | string[] | null | undefined;
}

interface SendMessagePayload extends WebSocketMessage {
  type: "new_message";
  content: string;
  reply_to_id?: string | null;
  reply_content?: string | null;
  image_url?: string | null;
}

interface MarkReadPayload extends WebSocketMessage {
  type: "mark_read";
}

interface NewMessageEvent extends WebSocketMessage {
  type: "new_message";
  message_id: string;
  content: string;
  sender_id: string;
  created_at: string;
  reply_to_id?: string | null;
  reply_content?: string | null;
  image_url?: string | null;
}

interface ReadReceiptEvent extends WebSocketMessage {
  type: "read_receipt";
  reader_username: string;
  read_at: string;
}

interface DeleteMessagePayload extends WebSocketMessage {
  type: "delete_message";
  message_id: string;
}

interface MessageDeletedEvent extends WebSocketMessage {
  type: "message_deleted";
  message_id: string;
}

export default useWebSocket;