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

const useCodeAssistant = ({ pageContainerRef }) => {
  // State variables for the assistant's UI and status
  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 [sessionId, setSessionId] = useState(null);

  // Refs for the WebSocket connection and directory handle
  const ws = useRef(null);
  const initialized = useRef(false);
  const dirHandleRef = useRef(null);

  // --- WebSocket Message Handlers ---
  const handleChatMessage = useCallback((message) => {
    setChatHistory((prev) => [...prev, { isUser: false, text: message.content }]);
    setIsProcessing(false);
  }, []);

  const handleThinkingLog = useCallback((message) => {
    setThinkingProcess((prev) => [
      ...prev,
      { type: message.subtype, message: message.content },
    ]);
  }, []);

  const handleError = useCallback((message) => {
    setErrorMessage(message.content);
    setShowErrorModal(true);
    setIsProcessing(false);
  }, []);

  const handleStatusUpdate = useCallback((message) => {
    setIsProcessing(message.processing);
    setIsPaused(message.paused);
    setConnectionStatus(message.status);
  }, []);

  const handleListDirectoryRequest = useCallback(async (message) => {
    const dirHandle = dirHandleRef.current;
    if (!dirHandle) {
      const errorMsg = "No folder selected by user.";
      console.warn(errorMsg);
      ws.current.send(JSON.stringify({ type: "error", content: errorMsg }));
      return;
    }
  
    try {
      const files = [];
  
      // Recursive function to walk through directories
      async function walkDirectory(handle, path = "") {
        for await (const [name, entry] of handle.entries()) {
          const relativePath = path ? `${path}/${name}` : name;
  
          if (entry.kind === "file") {
            files.push(relativePath); // Store full relative path
          } else if (entry.kind === "directory") {
            await walkDirectory(entry, relativePath); // Recurse into subdirectory
          }
        }
      }
  
      await walkDirectory(dirHandle);
  
      ws.current.send(
        JSON.stringify({
          type: "list_directory_response",
          files,
          request_id: message.request_id,
        })
      );
  
      setThinkingProcess((prev) => [
        ...prev,
        {
          type: "system",
          message: `Sent list of file names from folder "${dirHandle.name}" to server. Total files: ${files.length}`,
        },
      ]);
    } catch (error) {
      console.error("Failed to list directory:", error);
      ws.current.send(
        JSON.stringify({
          type: "error",
          content: "Failed to access folder contents.",
          request_id: message.request_id,
        })
      );
    }
  }, []);
  

  const handleReadFileRequest = useCallback(async (message) => {
    const dirHandle = dirHandleRef.current;
    const { filename, request_id } = message;

    if (!dirHandle) {
      ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id }));
      return;
    }

    try {
      const fileHandle = await dirHandle.getFileHandle(filename);
      const file = await fileHandle.getFile();
      const content = await file.text();

      ws.current.send(JSON.stringify({
        type: "file_content_response",
        filename,
        content,
        request_id,
      }));

      setThinkingProcess((prev) => [
        ...prev,
        { type: "system", message: `Sent content of file "${filename}" to server.` },
      ]);
    } catch (error) {
      console.error(`Failed to read file ${filename}:`, error);
      ws.current.send(JSON.stringify({
        type: "error",
        content: `Could not read file: ${filename}`,
        request_id,
      }));
    }
  }, []);

  const handleExecuteCommandRequest = useCallback((message) => {
    const { command, request_id } = message;

    const output = `Simulated output for command: '${command}'`;

    ws.current.send(JSON.stringify({
      type: "command_output",
      command,
      output,
      request_id,
    }));

    setThinkingProcess((prev) => [
      ...prev,
      { type: "system", message: `Simulated execution of command: '${command}'` },
    ]);
  }, []);

  // Main message handler that routes messages to the correct function
  const handleIncomingMessage = useCallback((message) => {
    switch (message.type) {
      case "chat_message":
        handleChatMessage(message);
        break;
      case "thinking_log":
        handleThinkingLog(message);
        break;
      case "error":
        handleError(message);
        break;
      case "status_update":
        handleStatusUpdate(message);
        break;
      case "list_directory":
        handleListDirectoryRequest(message);
        break;
      case "read_file":
        handleReadFileRequest(message);
        break;
      case "execute_command":
        handleExecuteCommandRequest(message);
        break;
      default:
        console.log("Unknown message type:", message);
    }
  }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFileRequest, handleExecuteCommandRequest]);

  // --- WebSocket Connection Setup ---
  useEffect(() => {
    if (initialized.current) return;
    initialized.current = true;

    const setupConnection = async () => {
      try {
        const { ws: newWs, sessionId: newSessionId } = await connectToWebSocket(
          handleIncomingMessage,
          () => setConnectionStatus("connected"),
          () => {
            setConnectionStatus("disconnected");
            setIsProcessing(false);
          },
          (error) => {
            setConnectionStatus("error");
            setErrorMessage(`Failed to connect: ${error.message}`);
            setShowErrorModal(true);
          }
        );
        ws.current = newWs;
        setSessionId(newSessionId);
      } catch (error) {
        console.error("Setup failed:", error);
      }
    };

    setupConnection();

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

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

  // Open folder picker and store handle
  const handleSelectFolder = useCallback(async (directoryHandle) => {
    if (!window.showDirectoryPicker) {
      // Don't use window.alert, use a custom modal
      return;
    }

    try {
      dirHandleRef.current = directoryHandle;
      setSelectedFolder(directoryHandle.name);

      ws.current.send(JSON.stringify({ type: "select_folder", path: directoryHandle.name }));

      setThinkingProcess((prev) => [
        ...prev,
        { type: "user", message: `Selected local folder: ${directoryHandle.name}` },
      ]);
    } catch (error) {
      console.error("Folder selection canceled or failed:", error);
    }
  }, []);

  // Control functions
  const handlePause = useCallback(() => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify({ type: "control", command: "pause" }));
    }
  }, []);

  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;