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);
  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);
    setConnectionStatus(message.status);
  }, []);

  const handleListDirectoryRequest = useCallback(async (message) => {
    const { request_id } = 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,
              created: file.created,
            });
          } 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: request_id,
          session_id: sessionId,
        })
      );

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

  // Helper function to recursively get a file handle from a full path
  const getFileHandleFromPath = async (dirHandle, filePath) => {
    // Split the path and filter out any empty parts to prevent the error
    const pathParts = filePath.split('/').filter(Boolean);
    let currentHandle = dirHandle;
    for (let i = 0; i < pathParts.length; i++) {
        const part = pathParts[i];
        try {
            if (i === pathParts.length - 1) {
                // Last part is the file name
                return await currentHandle.getFileHandle(part);
            } else {
                // Part is a directory
                currentHandle = await currentHandle.getDirectoryHandle(part);
            }
        } catch (error) {
            console.error(`Error navigating to path part '${part}':`, error);
            // Return null or re-throw based on desired error handling
            return null;
        }
    }
    return null;
  };


  const handleReadFilesRequest = useCallback(async (message) => {
    console.log(message);
    const { filepaths, request_id } = message;
    const dirHandle = dirHandleRef.current;
    if (!dirHandle) {
      ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id, session_id: sessionId }));
      return;
    }
  
    const filesData = [];
    const readFiles = []; // Array to store names of successfully read files
  
    for (const filepath of filepaths) {
      try {
        const fileHandle = await getFileHandleFromPath(dirHandle, filepath);
        const file = await fileHandle.getFile();
        const content = await file.text();
        filesData.push({ filepath, content });
        readFiles.push(filepath); // Add successfully read file to the list
      } catch (error) {
        console.error(`Failed to read file ${filepath}:`, error);
        ws.current.send(JSON.stringify({
          type: "error",
          content: `Could not read file: ${filepath}`,
          request_id: request_id,
          session_id: sessionId,
        }));
      }
    }
  
    ws.current.send(JSON.stringify({
      type: "file_content_response",
      files: filesData,
      request_id: request_id,
      session_id: sessionId,
    }));
  
    // Consolidate all thinking process updates into a single call
    setThinkingProcess((prev) => {
      const newMessages = [
        { type: "local", message: `Reading content of ${filepaths.length} files...` }
      ];
  
      // Add a message summarizing the files that were read
      if (readFiles.length > 0) {
        const displayMessage = readFiles.length > 10
          ? `Read ${readFiles.length} files successfully.`
          : `Read files successfully: [${readFiles.join(', ')}]`;
        newMessages.push({ type: "local", message: displayMessage });
      }
  
      newMessages.push({
        type: "local",
        message: `Sent content for ${filesData.length} files to server.`
      });
  
      return [...prev, ...newMessages];
    });
  }, [sessionId]);

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

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

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

  // 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(async (text) => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      setChatHistory((prev) => [...prev, { isUser: true, text }]);
      setIsProcessing(true);
      // Removed the extra call to getSessionId as it is already in state
      ws.current.send(JSON.stringify({ type: "chat_message",content: text, session_id: sessionId, path: dirHandleRef .current ? dirHandleRef.current.name : null  }));
    }
  }, [sessionId]);

  // 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, session_id: sessionId }));

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

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

  const handleStop = useCallback(() => {
    if (ws.current && ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(JSON.stringify({ type: "control", command: "stop", session_id: sessionId }));
    }
  }, [sessionId]);

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

export default useCodeAssistant;