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

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 [sessionId, setSessionId] = useState(null);

  const sessionIdRef = useRef(null); // ✅ Always current sessionId
  const ws = useRef(null);
  const initialized = useRef(false);
  const dirHandleRef = useRef(null);

  const handleChatMessage = useCallback((message) => {
    console.log("Received chat message:", message);
    setChatHistory((prev) => [...prev, {
      isUser: false,
      text: message.content,
      dicision: message.dicision,
      code_diff: message.code_diff,
      reasoning: message.reasoning
    }]);
    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;
    console.log("Received list directory request:", 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;
    }

    try {
      const files = [];

      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);
          }
        }
      }

      await walkDirectory(dirHandle);

      ws.current.send(JSON.stringify({
        type: "list_directory_response",
        files,
        request_id,
        session_id: sessionIdRef.current, // ✅ Always current
      }));

      setThinkingProcess((prev) => [...prev, {
        type: "local",
        message: `Sent ${files.length} files metadata information 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: sessionIdRef.current,
      }));
    }
  }, []);

  const getFileHandleFromPath = async (dirHandle, filePath) => {
    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) {
          return await currentHandle.getFileHandle(part);
        } else {
          currentHandle = await currentHandle.getDirectoryHandle(part);
        }
      } catch (error) {
        console.warn(`Path not found: ${filePath}`, error);
        return null;
      }
    }
    return null;
  };

  const handleReadFilesRequest = useCallback(async (message) => {
    const { filepaths, request_id } = message;
    console.log("Received read files request:", message);
    const dirHandle = dirHandleRef.current;

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

    const filesData = [];
    const readFiles = [];

    for (const filepath of filepaths) {
      try {
        const fileHandle = await getFileHandleFromPath(dirHandle, filepath);
        if (!fileHandle) {
          filesData.push({ filepath, content: null });
        }
        const file = await fileHandle.getFile();
        const content = await file.text();
        filesData.push({ filepath, content });
        readFiles.push(filepath);
      } catch (error) {
        console.warn(`Failed to read file: ${filepath}`, error);
        // ws.current.send(JSON.stringify({
        //   type: "error",
        //   content: `Could not read file: ${filepath}`,
        //   request_id,
        //   session_id: sessionIdRef.current,
        // }));
      }
    }

    ws.current.send(JSON.stringify({
      type: "file_content_response",
      files: filesData,
      request_id,
      session_id: sessionIdRef.current,
    }));

    setThinkingProcess((prev) => {
      const newMessages = [];

      if (readFiles.length > 0) {
        const displayMessage = readFiles.length > 10
          ? `Read ${readFiles.length} files successfully.`
          : `Read files and send successfully: [${readFiles.map(f => `"${f}"`).join(', ')}]`;
        newMessages.push({ type: "local", message: displayMessage });
      }
      return [...prev, ...newMessages];
    });
  }, []);

  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: sessionIdRef.current,
    }));

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

  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 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);
        sessionIdRef.current = newSessionId; // ✅ Keep ref in sync
      } catch (error) {
        console.error("Setup failed:", error);
      }
    };

    setupConnection();

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

  const handleSendChat = useCallback(async (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,
        session_id: sessionIdRef.current,
        path: dirHandleRef.current ? dirHandleRef.current.name : null
      }));
    }
  }, []);

  const handleSelectFolder = useCallback(async (directoryHandle) => {
    if (!window.showDirectoryPicker) return;

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

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

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

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

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

export default useCodeAssistant;