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

import { IMessage } from "../interfaces/conversation";
import { conversationWebSocket } from "../services/websockets/conversation";

interface StreamingMessage extends IMessage {
  isComplete?: boolean;
}

const useWebSocket = () => {
  const [messages, setMessages] = useState<StreamingMessage[]>([]);
  const [streamingMessage, setStreamingMessage] = useState("");
  const [isStreaming, setIsStreaming] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);

  const simulateStreaming = (completeMessage: string) => {
    setIsStreaming(true);
    setStreamingMessage("");
    const cleanMessage = completeMessage.replace(/<userStyle>.*<\/userStyle>/, "").trim();
    const chars = cleanMessage.split("");
    let currentIndex = -1;
    setStreamingMessage("");
    const streamInterval = setInterval(() => {
      if (currentIndex < chars.length) {
        currentIndex++;
        setStreamingMessage((prev) => prev + chars[currentIndex]);
      } else {
        clearInterval(streamInterval);
        setMessages((prev) => [
          ...prev,
          {
            text: cleanMessage,
            sender: "bot",
            isComplete: true,
          },
        ]);
        setIsStreaming(false);
        setStreamingMessage("");
      }
    }, 20);
    return () => clearInterval(streamInterval);
  };

  const connectWebSocket = useCallback((conversationId: string) => {
    const ws = conversationWebSocket({ conversationId, modality: "text" });
    wsRef.current = ws;

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

      if (!data.ai_message) return;

      simulateStreaming(data.ai_message);
    };

    ws.onclose = () => {
      wsRef.current = null;
    };

    ws.onerror = (error) => {
      console.error("WebSocket error:", error);
      wsRef.current = null;
    };
  }, []);

  useEffect(() => {
    return () => {
      wsRef.current?.close();
      wsRef.current = null;
    };
  }, []);

  const sendMessage = useCallback((message: string) => {
    if (message.trim()) {
      setMessages((prev) => [
        ...prev,
        {
          text: message,
          sender: "user",
          isComplete: true,
        },
      ]);
      if (wsRef.current?.readyState === WebSocket.OPEN) {
        wsRef.current.send(JSON.stringify({ user_message: message }));
      }
    }
  }, []);

  const getAllMessages = useCallback(() => {
    const displayMessages = [...messages];
    if (isStreaming) {
      displayMessages.push({
        text: streamingMessage,
        sender: "bot",
        isComplete: false,
      });
    }
    return displayMessages;
  }, [messages, isStreaming, streamingMessage]);

  return {
    messages: getAllMessages(),
    sendMessage,
    connectWebSocket,
    isStreaming,
  };
};

export default useWebSocket;