// src/hooks/useCodeAssistant.js
import { useState, useEffect, useRef, useCallback } from "react";

const useCodeAssistant = ({ pageContainerRef }) => {
  const [chatHistory, setChatHistory] = useState([]);
  const [thinkingProcess, setThinkingProcess] = useState([]);
  const [selectedFolder, setSelectedFolder] = useState(null);
  const [connectionStatus, setConnectionStatus] = useState("disconnected");
  const [isProcessing, setIsProcessing] = useState(false);
  const [isPaused, setIsPaused] = useState(false);
  const [errorMessage, setErrorMessage] = useState("");
  const [showErrorModal, setShowErrorModal] = useState(false);

  const ws = useRef(null);

  // Initialize WebSocket connection
  useEffect(() => {
    // Replace with your WebSocket server URL
    const websocketUrl = "wss://your-llm-assistant-server.com";
    ws.current = new WebSocket(websocketUrl);

    ws.current.onopen = () => {
      console.log("WebSocket connected");
      setConnectionStatus("connected");
    };

    ws.current.onmessage = (event) => {
      const message = JSON.parse(event.data);
      handleIncomingMessage(message);
    };

    ws.current.onclose = () => {
      console.log("WebSocket disconnected");
      setConnectionStatus("disconnected");
      setIsProcessing(false);
    };

    ws.current.onerror = (error) => {
      console.error("WebSocket error:", error);
      setConnectionStatus("error");
      setErrorMessage("WebSocket connection failed. Please check the server.");
      setShowErrorModal(true);
    };

    return () => {
      ws.current.close();
    };
  }, []);

  // Handle incoming messages from the server
  const handleIncomingMessage = useCallback((message) => {
    switch (message.type) {
      case "chat_message":
        setChatHistory((prevHistory) => [...prevHistory, { isUser: false, text: message.content }]);
        setIsProcessing(false);
        break;
      case "thinking_log":
        setThinkingProcess((prevProcess) => [
          ...prevProcess,
          { type: message.subtype, message: message.content },
        ]);
        break;
      case "file_lookup_request":
        // In a real app, this would trigger a local file read and send the content back.
        // For this hook, we just log the request.
        setThinkingProcess((prevProcess) => [
          ...prevProcess,
          { type: "system", message: `AI requested file: ${message.filename}` },
        ]);
        // To send a file back:
        // ws.current.send(JSON.stringify({ type: "file_content", content: "..." }));
        break;
      case "error":
        setErrorMessage(message.content);
        setShowErrorModal(true);
        setIsProcessing(false);
        break;
      case "status_update":
        setIsProcessing(message.processing);
        setIsPaused(message.paused);
        setConnectionStatus(message.status);
        break;
      default:
        console.log("Unknown message type:", message);
    }
  }, []);

  // Send a chat message to the server
  const handleSendChat = useCallback(
    (text) => {
      if (ws.current && ws.current.readyState === WebSocket.OPEN) {
        setChatHistory((prevHistory) => [...prevHistory, { isUser: true, text }]);
        setIsProcessing(true);
        ws.current.send(JSON.stringify({ type: "chat_message", content: text }));
      }
    },
    []
  );

  // Send the selected folder path to the server
  const handleSelectFolder = useCallback(
    (folderPath) => {
      if (ws.current && ws.current.readyState === WebSocket.OPEN) {
        setSelectedFolder(folderPath);
        ws.current.send(JSON.stringify({ type: "select_folder", path: folderPath }));
        setThinkingProcess((prevProcess) => [
          ...prevProcess,
          { type: "user", message: `Selected local folder: ${folderPath}` },
        ]);
      }
    },
    []
  );

  // Pause the AI's processing
  const handlePause = useCallback(() => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify({ type: "control", command: "pause" }));
    }
  }, []);

  // Stop the AI's processing
  const handleStop = useCallback(() => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify({ type: "control", command: "stop" }));
    }
  }, []);

  return {
    chatHistory,
    thinkingProcess,
    selectedFolder,
    connectionStatus,
    isProcessing,
    isPaused,
    errorMessage,
    showErrorModal,
    handleSendChat,
    handleSelectFolder,
    handlePause,
    handleStop,
    setShowErrorModal,
  };
};

export default useCodeAssistant;