import { useEffect, useRef, useState } from "react";
import { useAppStore } from "../store/useAppStore";
import { getWsUrl } from "@/config/api";

export function useWebSocket(url = "/ws/flow") {
  const [messages, setMessages] = useState([]);
  const [connected, setConnected] = useState(false);
  const wsRef = useRef(null);
  const setWsConnected = useAppStore((state) => state.setWsConnected);
  const addAlert = useAppStore((state) => state.addAlert);

  useEffect(() => {
    // Create WebSocket connection using configured WS URL
    const ws = new WebSocket(getWsUrl(url));
    wsRef.current = ws;

    ws.onopen = () => {
      console.log("WebSocket connected");
      setConnected(true);
      setWsConnected(true);

      // Subscribe to channels
      ws.send(
        JSON.stringify({
          type: "subscribe",
          channels: ["flow", "alerts"],
        })
      );
    };

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

        if (data.type === "data") {
          setMessages((prev) => [data, ...prev].slice(0, 100)); // Keep last 100

          // Handle alerts
          if (data.channel === "alerts" && data.data?.alert) {
            addAlert(data.data.alert);
          }
        }
      } catch (error) {
        console.error("Error parsing WebSocket message:", error);
      }
    };

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

    ws.onclose = () => {
      console.log("WebSocket disconnected");
      setConnected(false);
      setWsConnected(false);

      // Attempt to reconnect after 3 seconds
      setTimeout(() => {
        if (wsRef.current?.readyState === WebSocket.CLOSED) {
          // Will trigger useEffect again
        }
      }, 3000);
    };

    return () => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.close();
      }
    };
  }, [url, setWsConnected, addAlert]);

  const sendMessage = (message) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    }
  };

  return {
    messages,
    connected,
    sendMessage,
  };
}

