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

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); // Corrected this line
  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: "remote", message: message.content, round: message.round },
    ]);
  }, []);

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

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

  const handleListDirectoryRequest = useCallback(async (message) => {
    const { request_id, round } = 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, request_id }));
      return;
    }

    // setThinkingProcess((prev) => [
    //   ...prev,
    //   { type: "system", message: `Scanning directory...`, round },
    // ]);

    try {
      const files = [];

      // Recursive function to walk through directories
      async function walkDirectory(handle, path = '') {
        for await (const entry of handle.values()) {
          const entryPath = `${path}/${entry.name}`;
          if (entry.kind === "file") {
            const file = await entry.getFile();
            files.push({
              name: file.name,
              path: entryPath,
              size: file.size,
              lastModified: file.lastModified,
            });
          } else if (entry.kind === "directory") {
            await walkDirectory(entry, entryPath); // Recurse into subdirectory
          }
        }
      }

      await walkDirectory(dirHandle);

      ws.current.send(
        JSON.stringify({
          type: "list_directory_response",
          files,
          request_id,
          round,
        })
      );

      setThinkingProcess((prev) => [
        ...prev,
        {
          type: "local",
          message: `Sent list of files (${files.length}) to server.`,
          round,
        },
      ]);
    } catch (error) {
      console.error("Failed to list directory:", error);
      ws.current.send(
        JSON.stringify({
          type: "error",
          content: "Failed to access folder contents.",
          request_id,
          round,
        })
      );
    }
  }, []);

  const handleReadFilesRequest = useCallback(async (message) => {
    console.log(message);
    const { filenames, request_id, round } = message;
    const dirHandle = dirHandleRef.current;
    if (!dirHandle) {
      ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id, round }));
      return;
    }

    setThinkingProcess((prev) => [
      ...prev,
      { type: "local", message: `Reading content of ${filenames.length} files...`, round },
    ]);

    const filesData = [];
    for (const filename of filenames) {
      try {
        const fileHandle = await dirHandle.getFileHandle(filename);
        const file = await fileHandle.getFile();
        const content = await file.text();
        filesData.push({ filename, content });
        setThinkingProcess((prev) => [
          ...prev,
          { type: "local", message: `Read file: ${filename}`, round },
        ]);
      } 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,
          round,
        }));
      }
    }
    
    ws.current.send(JSON.stringify({
      type: "file_content_response",
      files: filesData,
      request_id,
      round,
    }));
    
    setThinkingProcess((prev) => [
      ...prev,
      { type: "local", message: `Sent content for ${filesData.length} files to server.`, round },
    ]);
  }, []);

  const handleExecuteCommandRequest = useCallback((message) => {
    const { command, request_id, round } = message;
    const output = `Simulated output for command: '${command}'`;

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

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

  // 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 "get_file_content":
        handleReadFilesRequest(message);
        break;
      case "execute_command":
        handleExecuteCommandRequest(message);
        break;
      default:
        console.log("Unknown message type:", message);
    }
  }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFilesRequest, 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) {
      return;
    }

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

      // Send the initial message to the server with a unique request_id
      const request_id = uuidv4();
      ws.current.send(JSON.stringify({ type: "select_folder_response", path: directoryHandle.name, request_id }));

      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;