dzone-front/src/chatbot/hooks/useSocketHandlers.ts

import { useCallback, useRef } from "react";
import { Socket } from "socket.io-client";
import { SOCKET_EVENTS } from "../constants";
import { Message } from "../types";

interface UseSocketHandlersProps {
  socket: Socket | null;
  setStatus: (status: string) => void;
  setIsTyping: (isTyping: boolean) => void;
  setIsProcessing: (isProcessing: boolean) => void;
  addMessage: (message: Partial<Message>) => Message;
  updateMessage: (id: string, updates: Partial<Message>) => void;
  addSystemMessage: (text: string, type?: string) => void;
  setCurrentSessionId: (sessionId: string | null) => void;
  isProcessing: boolean;
}

export const useSocketHandlers = ({
  socket,
  setStatus,
  setIsTyping,
  setIsProcessing,
  addMessage,
  updateMessage,
  addSystemMessage,
  setCurrentSessionId,
  isProcessing,
}: UseSocketHandlersProps) => {
  // Keep track of the current bot message like in the JS version
  const botMessageRef = useRef<{
    id: string;
    text: string;
  } | null>(null);

  const setupSocketHandlers = useCallback(() => {
    if (!socket) return () => {};

    // Connection events
    const onConnect = () => {
      console.log("🔌 Socket.IO connected!");
      setStatus("متصل شد ✅");
      addSystemMessage("اتصال برقرار شد", "success");
    };

    const onDisconnect = (reason: string) => {
      console.log("🔌 Socket.IO disconnected:", reason);
      setStatus("قطع ارتباط");
    };

    const onConnectionError = (error: any) => {
      console.error("❌ Socket.IO connection error:", error);
      addSystemMessage("خطا در اتصال", "error");
      setStatus("خطا در اتصال");
    };

    // Authentication events
    const onAuthenticated = (data: any) => {
      console.log("✅ Authenticated:", data);
      setStatus("احراز هویت شد ✅");
    };

    const onUnauthorized = (data: any) => {
      console.error("❌ Unauthorized:", data);
      addSystemMessage("خطا در احراز هویت", "error");
      setStatus("خطا در احراز هویت");
    };

    // Session events
    const onSessionCreated = (data: any) => {
      console.log("✅ Session created:", data);
      setCurrentSessionId(data.session.id);
      addSystemMessage("جلسه ایجاد شد 🎉", "success");
      setStatus("آماده برای گفتگو");

      // Auto-join the created session
      if (socket && data.session.id) {
        const userId =
          localStorage.getItem(import.meta.env.VITE_USER_ID_KEY || "user_id") ||
          localStorage.getItem("userId") ||
          "anonymous_user";

        socket.emit(SOCKET_EVENTS.JOIN_CHAT, {
          sessionId: data.session.id,
          userId: userId,
        });
      }
    };

    const onChatJoined = (data: any) => {
      console.log("✅ Joined chat:", data);
      addSystemMessage("به چت متصل شدید", "success");
      setStatus("متصل و آماده ✅");
    };

    const onSessionError = (data: any) => {
      console.error("❌ Session error:", data);
      addSystemMessage("خطا در ایجاد جلسه", "error");
    };

    // Message events
    const onMessageReceived = (data: any) => {
      console.log("📩 Message received:", data);
      // Message is already shown in UI when sent
    };

    // Bot response streaming - following the JS implementation pattern
    const onBotResponseStart = (data: any) => {
      console.log("🤖 Bot response starting:", data);

      // Turn off typing indicator since we're starting streaming
      setIsTyping(false);

      // Clear any existing bot message first
      if (botMessageRef.current) {
        console.log("⚠️ Clearing existing bot message before starting new one");
        botMessageRef.current = null;
      }

      // Create new bot message with streaming flag
      const newMessage = addMessage({
        text: "",
        isUser: false,
        streaming: true,
      });

      // Store reference like in JS version
      botMessageRef.current = {
        id: newMessage.id,
        text: "",
      };

      console.log("Created new streaming message with ID:", newMessage.id);
    };

    const onBotResponseChunk = (data: any) => {
      console.log("🤖 Bot response chunk received:", data);

      // If no bot message exists, create one (in case bot_response_start was missed)
      if (!botMessageRef.current && data?.data) {
        console.log("⚠️ Creating bot message from chunk (missing start event)");

        // Turn off typing indicator since we're starting streaming
        setIsTyping(false);

        const newMessage = addMessage({
          text: "",
          isUser: false,
          streaming: true,
        });

        botMessageRef.current = {
          id: newMessage.id,
          text: "",
        };
      }

      if (botMessageRef.current && data?.data) {
        // Accumulate text like in JS version
        botMessageRef.current.text += data.data;

        // Update the message content
        updateMessage(botMessageRef.current.id, {
          text: botMessageRef.current.text,
        });

        console.log(
          `Updated message to length: ${botMessageRef.current.text.length}`
        );
      } else {
        console.warn("Received chunk but no active bot message found", {
          hasBotMessage: !!botMessageRef.current,
          hasData: !!data?.data,
          data,
        });
      }
    };

    const onBotResponseEnd = (data: any) => {
      console.log("🤖 Bot response ended:", data);

      if (botMessageRef.current) {
        // Finalize the message by removing streaming flag
        updateMessage(botMessageRef.current.id, {
          streaming: false,
        });

        console.log(`Finalized message ${botMessageRef.current.id}`);

        // Clear the reference
        botMessageRef.current = null;
      }

      // Reset all states to allow user to send messages again
      setIsTyping(false);
      setIsProcessing(false);
      console.log("✅ Bot response ended - user can send messages again");
    };

    const onBotResponse = (data: any) => {
      console.log("🤖 Bot response (non-streaming):", data);

      if (data.message && data.message.content) {
        addMessage({
          text: data.message.content,
          isUser: false,
          metadata: data.message.metadata,
        });
      } else {
        console.warn("Received bot_response event with no content", data);
      }

      // Reset all states to allow user to send messages again
      setIsTyping(false);
      setIsProcessing(false);
      console.log("✅ Bot response completed - user can send messages again");
    };

    // Typing events
    const onTypingStart = (data: any) => {
      if (data.type === "bot") {
        setIsTyping(true);
        setStatus("در حال پاسخ...");
      }
    };

    const onTypingStop = (data: any) => {
      if (data.type === "bot") {
        setIsTyping(false);
        setStatus("آماده برای گفتگو");
      }
    };

    // User status events
    const onUserJoined = (data: any) => {
      console.log("👤 User joined:", data);
      addSystemMessage(`کاربر به چت پیوست`, "info");
    };

    const onUserLeft = (data: any) => {
      console.log("👤 User left:", data);
      addSystemMessage(`کاربر چت را ترک کرد`, "info");
    };

    // Error events
    const onError = (data: any) => {
      console.error("❌ Socket.IO error:", data);
      addSystemMessage(data.message || "خطا در پردازش", "error");

      // Reset states on error to allow user to try again
      setIsProcessing(false);
      setIsTyping(false);
      console.log("❌ Error occurred - reset processing state");
    };

    const onMessageError = (data: any) => {
      console.error("❌ Message error:", data);
      addSystemMessage(data.message || "خطا در ارسال پیام", "error");

      // Reset states on error to allow user to try again
      setIsProcessing(false);
      setIsTyping(false);
      console.log("❌ Message error - reset processing state");
    };

    // Register event handlers
    socket.on(SOCKET_EVENTS.CONNECT, onConnect);
    socket.on(SOCKET_EVENTS.DISCONNECT, onDisconnect);
    socket.on(SOCKET_EVENTS.CONNECTION_ERROR, onConnectionError);
    socket.on(SOCKET_EVENTS.AUTHENTICATED, onAuthenticated);
    socket.on(SOCKET_EVENTS.UNAUTHORIZED, onUnauthorized);
    socket.on(SOCKET_EVENTS.SESSION_CREATED, onSessionCreated);
    socket.on(SOCKET_EVENTS.CHAT_JOINED, onChatJoined);
    socket.on(SOCKET_EVENTS.SESSION_ERROR, onSessionError);
    socket.on(SOCKET_EVENTS.MESSAGE_RECEIVED, onMessageReceived);
    socket.on(SOCKET_EVENTS.BOT_RESPONSE_START, onBotResponseStart);
    socket.on(SOCKET_EVENTS.BOT_RESPONSE_CHUNK, onBotResponseChunk);
    socket.on(SOCKET_EVENTS.BOT_RESPONSE_END, onBotResponseEnd);
    socket.on(SOCKET_EVENTS.BOT_RESPONSE, onBotResponse);
    socket.on(SOCKET_EVENTS.TYPING_START, onTypingStart);
    socket.on(SOCKET_EVENTS.TYPING_STOP, onTypingStop);
    socket.on(SOCKET_EVENTS.USER_JOINED, onUserJoined);
    socket.on(SOCKET_EVENTS.USER_LEFT, onUserLeft);
    socket.on(SOCKET_EVENTS.ERROR, onError);
    socket.on(SOCKET_EVENTS.MESSAGE_ERROR, onMessageError);
    socket.on(SOCKET_EVENTS.WEBSOCKET_ERROR, onError);

    // Return cleanup function
    return () => {
      socket.off(SOCKET_EVENTS.CONNECT, onConnect);
      socket.off(SOCKET_EVENTS.DISCONNECT, onDisconnect);
      socket.off(SOCKET_EVENTS.CONNECTION_ERROR, onConnectionError);
      socket.off(SOCKET_EVENTS.AUTHENTICATED, onAuthenticated);
      socket.off(SOCKET_EVENTS.UNAUTHORIZED, onUnauthorized);
      socket.off(SOCKET_EVENTS.SESSION_CREATED, onSessionCreated);
      socket.off(SOCKET_EVENTS.CHAT_JOINED, onChatJoined);
      socket.off(SOCKET_EVENTS.SESSION_ERROR, onSessionError);
      socket.off(SOCKET_EVENTS.MESSAGE_RECEIVED, onMessageReceived);
      socket.off(SOCKET_EVENTS.BOT_RESPONSE_START, onBotResponseStart);
      socket.off(SOCKET_EVENTS.BOT_RESPONSE_CHUNK, onBotResponseChunk);
      socket.off(SOCKET_EVENTS.BOT_RESPONSE_END, onBotResponseEnd);
      socket.off(SOCKET_EVENTS.BOT_RESPONSE, onBotResponse);
      socket.off(SOCKET_EVENTS.TYPING_START, onTypingStart);
      socket.off(SOCKET_EVENTS.TYPING_STOP, onTypingStop);
      socket.off(SOCKET_EVENTS.USER_JOINED, onUserJoined);
      socket.off(SOCKET_EVENTS.USER_LEFT, onUserLeft);
      socket.off(SOCKET_EVENTS.ERROR, onError);
      socket.off(SOCKET_EVENTS.MESSAGE_ERROR, onMessageError);
      socket.off(SOCKET_EVENTS.WEBSOCKET_ERROR, onError);

      // Clear bot message reference on cleanup
      botMessageRef.current = null;
    };
  }, [
    socket,
    setStatus,
    setIsTyping,
    setIsProcessing,
    addMessage,
    updateMessage,
    addSystemMessage,
    setCurrentSessionId,
  ]);

  return { setupSocketHandlers };
};

// Helper function to format text with markdown-like syntax
const formatText = (text: string) => {
  if (!text) return "";

  // Clean the text first to avoid HTML injection
  const cleanedText = text
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;")
    .replace(/'/g, "&#039;");

  // Format with markdown-like syntax and handle newlines properly
  const formattedText = cleanedText
    // Bold text
    .replace(/\*\*(.*?)\*\*/g, "<strong>$1</strong>")
    // Italic text
    .replace(/\*(.*?)\*/g, "<em>$1</em>")
    // Line breaks - ensure they work properly with proper spacing
    .replace(/\n/g, "<br class='message-line-break' />")
    // Links (optional)
    .replace(
      /\[([^\]]+)\]\(([^)]+)\)/g,
      '<a href="$2" target="_blank" rel="noopener noreferrer" class="text-blue-500 underline">$1</a>'
    );

  return formattedText;
};