"use client";
import React, {
  createContext,
  useState,
  useCallback,
  useEffect,
  useMemo,
  useRef,
  MouseEvent,
  Suspense,
} from "react";
import Link from "next/link";
import {
  ReactFlow,
  ReactFlowProvider,
  Controls,
  Background,
  Node,
  OnConnect,
  Connection,
  MarkerType,
  NodeChange,
  EdgeChange,
  useReactFlow,
  applyEdgeChanges,
  applyNodeChanges,
} from "@xyflow/react";
import "@xyflow/react/dist/style.css";
import { ConnectedEdge, CustomNode } from "../CustomNode/CustomNode";
import "./flow.css";
import {
  BlockIORootSchema,
  BlockUIType,
  formatEdgeID,
  GraphExecutionID,
  GraphID,
  GraphMeta,
  LibraryAgent,
  SpecialBlockID,
} from "@/lib/autogpt-server-api";
import { getV1GetSpecificGraph } from "@/app/api/__generated__/endpoints/graphs/graphs";
import { okData } from "@/app/api/helpers";
import { IncompatibilityInfo } from "../../../hooks/useSubAgentUpdate/types";
import { Key, storage } from "@/services/storage/local-storage";
import { findNewlyAddedBlockCoordinates, getTypeColor } from "@/lib/utils";
import { history } from "../history";
import { CustomEdge } from "../CustomEdge/CustomEdge";
import ConnectionLine from "../ConnectionLine";
import {
  Control,
  ControlPanel,
} from "@/app/(platform)/build/components/legacy-builder/ControlPanel";
import { SaveControl } from "@/app/(platform)/build/components/legacy-builder/SaveControl";
import { BlocksControl } from "@/app/(platform)/build/components/legacy-builder/BlocksControl";
import { GraphSearchControl } from "@/app/(platform)/build/components/legacy-builder/GraphSearchControl";
import { IconUndo2, IconRedo2 } from "@/components/__legacy__/ui/icons";
import {
  Alert,
  AlertDescription,
  AlertTitle,
} from "@/components/molecules/Alert/Alert";
import { startTutorial } from "../tutorial";
import useAgentGraph from "@/hooks/useAgentGraph";
import { v4 as uuidv4 } from "uuid";
import { useRouter, usePathname, useSearchParams } from "next/navigation";
import RunnerUIWrapper, { RunnerUIWrapperRef } from "../RunnerUIWrapper";
import OttoChatWidget from "@/app/(platform)/build/components/legacy-builder/OttoChatWidget";
import { useToast } from "@/components/molecules/Toast/use-toast";
import { useCopyPaste } from "../useCopyPaste";
import NewControlPanel from "@/app/(platform)/build/components/NewControlPanel/NewControlPanel";
import { Flag, useGetFlag } from "@/services/feature-flags/use-get-flag";
import { BuildActionBar } from "../BuildActionBar";
import { FloatingReviewsPanel } from "@/components/organisms/FloatingReviewsPanel/FloatingReviewsPanel";
import { useFlowRealtime } from "@/app/(platform)/build/components/FlowEditor/Flow/useFlowRealtime";
import { FloatingSafeModeToggle } from "../../FloatingSafeModeToogle";

// This is for the history, this is the minimum distance a block must move before it is logged
// It helps to prevent spamming the history with small movements especially when pressing on a input in a block
const MINIMUM_MOVE_BEFORE_LOG = 50;

export type ResolutionModeState = {
  active: boolean;
  nodeId: string | null;
  incompatibilities: IncompatibilityInfo | null;
  brokenEdgeIds: string[];
  pendingUpdate: Record<string, unknown> | null; // The hardcoded values to apply after resolution
};

type BuilderContextType = {
  libraryAgent: LibraryAgent | null;
  visualizeBeads: "no" | "static" | "animate";
  setIsAnyModalOpen: (isOpen: boolean) => void;
  getNextNodeId: () => string;
  getNodeTitle: (nodeID: string) => string | null;
  availableFlows: GraphMeta[];
  resolutionMode: ResolutionModeState;
  enterResolutionMode: (
    nodeId: string,
    incompatibilities: IncompatibilityInfo,
    brokenEdgeIds: string[],
    pendingUpdate: Record<string, unknown>,
  ) => void;
  exitResolutionMode: () => void;
  applyPendingUpdate: () => void;
};

export type NodeDimension = {
  [nodeId: string]: {
    x: number;
    y: number;
    width: number;
    height: number;
  };
};

export const BuilderContext = createContext<BuilderContextType | null>(null);

const FlowEditor: React.FC<{
  flowID?: GraphID;
  flowVersion?: number;
  className?: string;
}> = ({ flowID, flowVersion, className }) => {
  const {
    addNodes,
    addEdges,
    getNode,
    deleteElements,
    updateNode,
    getViewport,
    setViewport,
    fitView,
    screenToFlowPosition,
  } = useReactFlow<CustomNode, CustomEdge>();
  const [nodeId, setNodeId] = useState<number>(1);
  const [isAnyModalOpen, setIsAnyModalOpen] = useState(false);
  const [visualizeBeads] = useState<"no" | "static" | "animate">("animate");
  const [flowExecutionID, setFlowExecutionID] = useState<
    GraphExecutionID | undefined
  >();
  // State to control if blocks menu should be pinned open
  const [pinBlocksPopover, setPinBlocksPopover] = useState(false);
  // State to control if save popover should be pinned open
  const [pinSavePopover, setPinSavePopover] = useState(false);
  const [hasAutoFramed, setHasAutoFramed] = useState(false);

  const {
    agentName,
    setAgentName,
    agentDescription,
    setAgentDescription,
    agentRecommendedScheduleCron,
    setAgentRecommendedScheduleCron,
    savedAgent,
    libraryAgent,
    availableBlocks,
    availableFlows,
    getOutputType,
    saveAgent,
    saveAndRun,
    stopRun,
    createRunSchedule,
    isSaving,
    isRunning,
    isStopping,
    isScheduling,
    graphExecutionError,
    nodes,
    setNodes,
    edges,
    setEdges,
  } = useAgentGraph(
    flowID,
    flowVersion,
    flowExecutionID,
    visualizeBeads !== "no",
  );
  const [immediateNodePositions, setImmediateNodePositions] = useState<
    Record<string, { x: number; y: number }>
  >(Object.fromEntries(nodes.map((node) => [node.id, node.position])));

  // Add realtime execution status tracking for FloatingReviewsPanel
  useFlowRealtime();

  const router = useRouter();
  const pathname = usePathname();
  const params = useSearchParams();
  const initialPositionRef = useRef<{
    [key: string]: { x: number; y: number };
  }>({});
  const isDragging = useRef(false);

  const runnerUIRef = useRef<RunnerUIWrapperRef>(null);

  const { toast } = useToast();

  // It stores the dimension of all nodes with position as well
  const [nodeDimensions, setNodeDimensions] = useState<NodeDimension>({});

  // Resolution mode state for sub-agent incompatible updates
  const [resolutionMode, setResolutionMode] = useState<ResolutionModeState>({
    active: false,
    nodeId: null,
    incompatibilities: null,
    brokenEdgeIds: [],
    pendingUpdate: null,
  });

  const enterResolutionMode = useCallback(
    (
      nodeId: string,
      incompatibilities: IncompatibilityInfo,
      brokenEdgeIds: string[],
      pendingUpdate: Record<string, unknown>,
    ) => {
      setResolutionMode({
        active: true,
        nodeId,
        incompatibilities,
        brokenEdgeIds,
        pendingUpdate,
      });
    },
    [],
  );

  const exitResolutionMode = useCallback(() => {
    setResolutionMode({
      active: false,
      nodeId: null,
      incompatibilities: null,
      brokenEdgeIds: [],
      pendingUpdate: null,
    });
  }, []);

  // Apply pending update after resolution mode completes
  const applyPendingUpdate = useCallback(() => {
    if (!resolutionMode.nodeId || !resolutionMode.pendingUpdate) return;

    const node = nodes.find((n) => n.id === resolutionMode.nodeId);
    if (node) {
      const pendingUpdate = resolutionMode.pendingUpdate as {
        [key: string]: any;
      };
      setNodes((nds) =>
        nds.map((n) =>
          n.id === resolutionMode.nodeId
            ? { ...n, data: { ...n.data, hardcodedValues: pendingUpdate } }
            : n,
        ),
      );
    }
    exitResolutionMode();
    toast({
      title: "Update complete",
      description: "Agent has been updated to the new version.",
    });
  }, [resolutionMode, nodes, setNodes, exitResolutionMode, toast]);

  // Check if all broken edges have been removed and auto-apply pending update
  useEffect(() => {
    if (!resolutionMode.active || resolutionMode.brokenEdgeIds.length === 0) {
      return;
    }

    const currentEdgeIds = new Set(edges.map((e) => e.id));
    const remainingBrokenEdges = resolutionMode.brokenEdgeIds.filter((id) =>
      currentEdgeIds.has(id),
    );

    if (remainingBrokenEdges.length === 0) {
      // All broken edges have been removed, apply pending update
      applyPendingUpdate();
    } else if (
      remainingBrokenEdges.length !== resolutionMode.brokenEdgeIds.length
    ) {
      // Update the list of broken edges
      setResolutionMode((prev) => ({
        ...prev,
        brokenEdgeIds: remainingBrokenEdges,
      }));
    }
  }, [edges, resolutionMode, applyPendingUpdate]);

  // Set page title with or without graph name
  useEffect(() => {
    document.title = savedAgent
      ? `${savedAgent.name} - Builder - AutoGPT Platform`
      : `Builder - AutoGPT Platform`;
  }, [savedAgent]);

  const graphHasWebhookNodes = useMemo(
    () =>
      nodes.some((n) =>
        [BlockUIType.WEBHOOK, BlockUIType.WEBHOOK_MANUAL].includes(
          n.data.uiType,
        ),
      ),
    [nodes],
  );

  useEffect(() => {
    if (params.get("resetTutorial") === "true") {
      storage.clean(Key.SHEPHERD_TOUR);
      router.push(pathname);
    } else if (!storage.get(Key.SHEPHERD_TOUR)) {
      const emptyNodes = (forceRemove: boolean = false) =>
        forceRemove ? (setNodes([]), setEdges([]), true) : nodes.length === 0;
      startTutorial(emptyNodes, setPinBlocksPopover, setPinSavePopover);
      storage.set(Key.SHEPHERD_TOUR, "yes");
    }
  }, [router, pathname, params, setEdges, setNodes, nodes.length]);

  useEffect(() => {
    if (params.get("open_scheduling") === "true") {
      runnerUIRef.current?.openRunInputDialog();
    }
    setFlowExecutionID(
      (params.get("flowExecutionID") as GraphExecutionID) || undefined,
    );
  }, [params]);

  useEffect(() => {
    const handleKeyDown = (event: KeyboardEvent) => {
      const isMac = navigator.platform.toUpperCase().indexOf("MAC") >= 0;
      const isUndo =
        (isMac ? event.metaKey : event.ctrlKey) && event.key === "z";
      const isRedo =
        (isMac ? event.metaKey : event.ctrlKey) &&
        (event.key === "y" || (event.shiftKey && event.key === "Z"));

      if (isUndo) {
        event.preventDefault();
        history.undo();
      }

      if (isRedo) {
        event.preventDefault();
        history.redo();
      }
    };

    window.addEventListener("keydown", handleKeyDown);

    return () => {
      window.removeEventListener("keydown", handleKeyDown);
    };
  }, []);

  const onNodeDragStart = (_: MouseEvent, node: Node) => {
    initialPositionRef.current[node.id] = { ...node.position };
    isDragging.current = true;
  };

  const onNodeDragEnd = (_: MouseEvent, node: Node | null) => {
    if (!node) return;

    isDragging.current = false;
    const oldPosition = initialPositionRef.current[node.id];
    const newPosition = node.position;

    // Clear immediate position, because on drag end it is no longer needed
    setImmediateNodePositions((prevPositions) => {
      const updatedPositions = { ...prevPositions };
      delete updatedPositions[node.id];
      return updatedPositions;
    });

    // Calculate the movement distance
    if (!oldPosition || !newPosition) return;

    const distanceMoved = Math.sqrt(
      Math.pow(newPosition.x - oldPosition.x, 2) +
        Math.pow(newPosition.y - oldPosition.y, 2),
    );

    if (distanceMoved > MINIMUM_MOVE_BEFORE_LOG) {
      // Minimum movement threshold
      history.push({
        type: "UPDATE_NODE_POSITION",
        payload: { nodeId: node.id, oldPosition, newPosition },
        undo: () => updateNode(node.id, { position: oldPosition }),
        redo: () => updateNode(node.id, { position: newPosition }),
      });
    }
    delete initialPositionRef.current[node.id];
  };

  // Function to clear status, output, and close the output info dropdown of all nodes
  // and reset data beads on edges
  const clearNodesStatusAndOutput = useCallback(() => {
    setNodes((nds) =>
      nds.map((node) => ({
        ...node,
        data: {
          ...node.data,
          status: undefined,
          isOutputOpen: false,
        },
      })),
    );
  }, [setNodes]);

  const onNodesChange = useCallback(
    (nodeChanges: NodeChange<CustomNode>[]) => {
      // Intercept position changes to update immediate positions & prevent excessive node re-renders
      const draggingPosChanges = nodeChanges
        .filter((c) => c.type === "position")
        .filter((c) => c.dragging === true);
      if (draggingPosChanges.length > 0) {
        setImmediateNodePositions((prevPositions) => {
          const newPositions = { ...prevPositions };
          draggingPosChanges.forEach((change) => {
            if (change.position) newPositions[change.id] = change.position;
          });
          return newPositions;
        });

        // Don't further process ongoing position changes
        nodeChanges = nodeChanges.filter(
          (change) => change.type !== "position" || change.dragging !== true,
        );
        if (nodeChanges.length === 0) return;
      }

      // Persist the changes
      setNodes((prev) => applyNodeChanges(nodeChanges, prev));

      // Remove all edges that were connected to deleted nodes
      nodeChanges
        .filter((change) => change.type === "remove")
        .forEach((deletedNode) => {
          const nodeID = deletedNode.id;
          const deletedNodeData = nodes.find((node) => node.id === nodeID);

          if (deletedNodeData) {
            history.push({
              type: "DELETE_NODE",
              payload: { node: deletedNodeData.data },
              undo: () => addNodes(deletedNodeData),
              redo: () => deleteElements({ nodes: [{ id: nodeID }] }),
            });
          }

          const connectedEdges = edges.filter((edge) =>
            [edge.source, edge.target].includes(nodeID),
          );
          deleteElements({
            edges: connectedEdges.map((edge) => ({ id: edge.id })),
          });
        });
    },
    [deleteElements, setNodes, nodes, edges, addNodes],
  );

  const onConnect: OnConnect = useCallback(
    (connection: Connection) => {
      // Check if this exact connection already exists
      const existingConnection = edges.find(
        (edge) =>
          edge.source === connection.source &&
          edge.target === connection.target &&
          edge.sourceHandle === connection.sourceHandle &&
          edge.targetHandle === connection.targetHandle,
      );

      if (existingConnection) {
        console.warn("This exact connection already exists.");
        return;
      }

      const edgeColor = getTypeColor(
        getOutputType(nodes, connection.source!, connection.sourceHandle!),
      );
      const sourceNode = getNode(connection.source!);
      const newEdge: CustomEdge = {
        id: formatEdgeID(connection),
        type: "custom",
        markerEnd: {
          type: MarkerType.ArrowClosed,
          strokeWidth: 2,
          color: edgeColor,
        },
        data: {
          edgeColor,
          sourcePos: sourceNode!.position,
          isStatic: sourceNode!.data.isOutputStatic,
          beadUp: 0,
          beadDown: 0,
        },
        ...connection,
        source: connection.source!,
        target: connection.target!,
      };

      addEdges(newEdge);
      history.push({
        type: "ADD_EDGE",
        payload: { edge: newEdge },
        undo: () => {
          deleteElements({ edges: [{ id: newEdge.id }] });
        },
        redo: () => {
          addEdges(newEdge);
        },
      });
      clearNodesStatusAndOutput(); // Clear status and output on connection change
    },
    [
      getNode,
      addEdges,
      deleteElements,
      clearNodesStatusAndOutput,
      nodes,
      edges,
      formatEdgeID,
      getOutputType,
    ],
  );

  const onEdgesChange = useCallback(
    (edgeChanges: EdgeChange<CustomEdge>[]) => {
      // Persist the changes
      setEdges((prev) => applyEdgeChanges(edgeChanges, prev));

      // Propagate edge changes to node data
      const addedEdges = edgeChanges.filter((change) => change.type === "add"),
        replaceEdges = edgeChanges.filter(
          (change) => change.type === "replace",
        ),
        removedEdges = edgeChanges.filter((change) => change.type === "remove");

      if (addedEdges.length > 0 || removedEdges.length > 0) {
        setNodes((nds) => {
          const newNodes = nds.map((node) => ({
            ...node,
            data: {
              ...node.data,
              connections: [
                // Remove node connections for deleted edges
                ...node.data.connections.filter(
                  (conn) =>
                    !removedEdges.some(
                      (removedEdge) => removedEdge.id === conn.id,
                    ),
                ),
                // Add node connections for added edges
                ...addedEdges.map(
                  (addedEdge): ConnectedEdge => ({
                    id: addedEdge.item.id,
                    source: addedEdge.item.source,
                    target: addedEdge.item.target,
                    sourceHandle: addedEdge.item.sourceHandle!,
                    targetHandle: addedEdge.item.targetHandle!,
                  }),
                ),
              ],
            },
          }));

          return newNodes;
        });

        if (removedEdges.length > 0) {
          clearNodesStatusAndOutput(); // Clear status and output on edge deletion
        }
      }

      if (replaceEdges.length > 0) {
        // Reset node connections for all edges
        console.warn(
          "useReactFlow().setRootEdges was used to overwrite all edges. " +
            "Use addEdges, deleteElements, or reconnectEdge for incremental changes.",
          replaceEdges,
        );
        setNodes((nds) =>
          nds.map((node) => ({
            ...node,
            data: {
              ...node.data,
              connections: [
                ...replaceEdges.map(
                  (replaceEdge): ConnectedEdge => ({
                    id: replaceEdge.item.id,
                    source: replaceEdge.item.source,
                    target: replaceEdge.item.target,
                    sourceHandle: replaceEdge.item.sourceHandle!,
                    targetHandle: replaceEdge.item.targetHandle!,
                  }),
                ),
              ],
            },
          })),
        );
        clearNodesStatusAndOutput();
      }
    },
    [setNodes, clearNodesStatusAndOutput, setEdges],
  );

  const getNextNodeId = useCallback(() => {
    return uuidv4();
  }, []);

  useEffect(() => {
    if (nodes.length === 0) {
      return;
    }

    if (hasAutoFramed) {
      return;
    }

    const rafId = requestAnimationFrame(() => {
      fitView({ padding: 0.2, duration: 800, maxZoom: 1 });
      setHasAutoFramed(true);
    });

    return () => cancelAnimationFrame(rafId);
  }, [fitView, hasAutoFramed, nodes.length]);

  useEffect(() => {
    setHasAutoFramed(false);
  }, [flowID, flowVersion]);

  const navigateToNode = useCallback(
    (nodeId: string) => {
      const node = getNode(nodeId);
      if (!node) return;

      // Center the viewport on the selected node
      const zoom = 1.2; // Slightly zoom in for better visibility
      const nodeX = node.position.x + (node.width || 500) / 2;
      const nodeY = node.position.y + (node.height || 400) / 2;

      setViewport({
        x: window.innerWidth / 2 - nodeX * zoom,
        y: window.innerHeight / 2 - nodeY * zoom,
        zoom: zoom,
      });

      // Add a temporary highlight effect to the node
      updateNode(nodeId, {
        style: {
          ...node.style,
          boxShadow: "0 0 20px 5px rgba(59, 130, 246, 0.8)",
          transition: "box-shadow 0.3s ease-in-out",
        },
      });

      // Remove highlight after a delay
      setTimeout(() => {
        updateNode(nodeId, {
          style: {
            ...node.style,
            boxShadow: undefined,
          },
        });
      }, 2000);
    },
    [getNode, setViewport, updateNode],
  );

  const highlightNode = useCallback(
    (nodeId: string | null) => {
      if (!nodeId) {
        // Clear all highlights
        nodes.forEach((node) => {
          updateNode(node.id, {
            style: {
              ...node.style,
              boxShadow: undefined,
            },
          });
        });
        return;
      }

      const node = getNode(nodeId);
      if (!node) return;

      // Add highlight effect without moving view
      updateNode(nodeId, {
        style: {
          ...node.style,
          boxShadow: "0 0 15px 3px rgba(59, 130, 246, 0.6)",
          transition: "box-shadow 0.2s ease-in-out",
        },
      });
    },
    [getNode, updateNode, nodes],
  );

  /* Shared helper to create and add a node */
  const createAndAddNode = useCallback(
    async (
      blockID: string,
      blockName: string,
      hardcodedValues: Record<string, any>,
      position: { x: number; y: number },
    ): Promise<CustomNode | null> => {
      const nodeSchema = availableBlocks.find((node) => node.id === blockID);
      if (!nodeSchema) {
        console.error(`Schema not found for block ID: ${blockID}`);
        return null;
      }

      // For agent blocks, fetch the full graph to get schemas
      let inputSchema: BlockIORootSchema = nodeSchema.inputSchema;
      let outputSchema: BlockIORootSchema = nodeSchema.outputSchema;
      let finalHardcodedValues = hardcodedValues;

      if (blockID === SpecialBlockID.AGENT) {
        const graphID = hardcodedValues.graph_id as string;
        const graphVersion = hardcodedValues.graph_version as number;
        const graphData = okData(
          await getV1GetSpecificGraph(graphID, { version: graphVersion }),
        );

        if (graphData) {
          inputSchema = graphData.input_schema as BlockIORootSchema;
          outputSchema = graphData.output_schema as BlockIORootSchema;
          finalHardcodedValues = {
            ...hardcodedValues,
            input_schema: graphData.input_schema,
            output_schema: graphData.output_schema,
          };
        } else {
          console.error("Failed to fetch graph data for agent block");
        }
      }

      const newNode: CustomNode = {
        id: nodeId.toString(),
        type: "custom",
        position,
        data: {
          blockType: blockName,
          blockCosts: nodeSchema.costs || [],
          title: `${blockName} ${nodeId}`,
          description: nodeSchema.description,
          categories: nodeSchema.categories,
          inputSchema: inputSchema,
          outputSchema: outputSchema,
          hardcodedValues: finalHardcodedValues,
          connections: [],
          isOutputOpen: false,
          block_id: blockID,
          isOutputStatic: nodeSchema.staticOutput,
          uiType: nodeSchema.uiType,
        },
      };

      addNodes(newNode);
      setNodeId((prevId) => prevId + 1);
      clearNodesStatusAndOutput();

      history.push({
        type: "ADD_NODE",
        payload: { node: { ...newNode, ...newNode.data } },
        undo: () => deleteElements({ nodes: [{ id: newNode.id }] }),
        redo: () => addNodes(newNode),
      });

      return newNode;
    },
    [
      availableBlocks,
      nodeId,
      addNodes,
      deleteElements,
      clearNodesStatusAndOutput,
    ],
  );

  const addNode = useCallback(
    async (
      blockId: string,
      nodeType: string,
      hardcodedValues: Record<string, any> = {},
    ) => {
      const nodeSchema = availableBlocks.find((node) => node.id === blockId);
      if (!nodeSchema) {
        console.error(`Schema not found for block ID: ${blockId}`);
        return;
      }

      /*
       Calculate a position to the right of the newly added block, allowing for some margin.
       If adding to the right side causes the new block to collide with an existing block, attempt to place it at the bottom or left.
       Why not the top? Because the height of the new block is unknown.
       If it still collides, run a loop to find the best position where it does not collide.
       Then, adjust the canvas to center on the newly added block.
       Note: The width is known, e.g., w = 300px for a note and w = 500px for others, but the height is dynamic.
       */

      // Alternative: We could also use D3 force, Intersection for this (React flow Pro examples)

      const { x, y } = getViewport();
      const position =
        nodeDimensions && Object.keys(nodeDimensions).length > 0
          ? findNewlyAddedBlockCoordinates(
              nodeDimensions,
              nodeSchema.uiType == BlockUIType.NOTE ? 300 : 500,
              60,
              1.0,
            )
          : {
              x: window.innerWidth / 2 - x,
              y: window.innerHeight / 2 - y,
            };

      const newNode = await createAndAddNode(
        blockId,
        nodeType,
        hardcodedValues,
        position,
      );
      if (!newNode) return;

      setViewport(
        {
          x: -position.x * 0.8 + (window.innerWidth - 0.0) / 2,
          y: -position.y * 0.8 + (window.innerHeight - 400) / 2,
          zoom: 0.8,
        },
        { duration: 500 },
      );
    },
    [
      getViewport,
      setViewport,
      availableBlocks,
      nodeDimensions,
      createAndAddNode,
    ],
  );

  const findNodeDimensions = useCallback(() => {
    const newNodeDimensions: NodeDimension = nodes.reduce((acc, node) => {
      const nodeElement = document.querySelector(
        `[data-id="custom-node-${node.id}"]`,
      );
      if (nodeElement) {
        const rect = nodeElement.getBoundingClientRect();
        const { left, top, width, height } = rect;

        const { x, y, zoom } = getViewport();

        // Convert screen coordinates to flow coordinates
        const flowX = (left - x) / zoom;
        const flowY = (top - y) / zoom;
        const flowWidth = width / zoom;
        const flowHeight = height / zoom;

        acc[node.id] = {
          x: flowX,
          y: flowY,
          width: flowWidth,
          height: flowHeight,
        };
      }
      return acc;
    }, {} as NodeDimension);

    setNodeDimensions(newNodeDimensions);
  }, [nodes, getViewport]);

  useEffect(() => {
    findNodeDimensions();
  }, [nodes, findNodeDimensions]);

  const getNodeTitle = useCallback(
    (nodeID: string) => {
      const node = nodes.find((n) => n.data.backend_id === nodeID);
      if (!node) return null;

      return (
        node.data.metadata?.customized_name ||
        (node.data.uiType == BlockUIType.AGENT &&
          node.data.hardcodedValues.agent_name) ||
        node.data.blockType.replace(/Block$/, "")
      );
    },
    [nodes],
  );

  const handleCopyPaste = useCopyPaste(getNextNodeId);

  const handleKeyDown = useCallback(
    (event: KeyboardEvent) => {
      // Prevent copy/paste if any modal is open or if the focus is on an input element
      const activeElement = document.activeElement;
      const isInputField =
        activeElement?.tagName === "INPUT" ||
        activeElement?.tagName === "TEXTAREA" ||
        activeElement?.getAttribute("contenteditable") === "true";

      if (isAnyModalOpen || isInputField) return;

      handleCopyPaste(event);
    },
    [isAnyModalOpen, handleCopyPaste],
  );

  useEffect(() => {
    window.addEventListener("keydown", handleKeyDown);
    return () => {
      window.removeEventListener("keydown", handleKeyDown);
    };
  }, [handleKeyDown]);

  const onNodesDelete = useCallback(() => {
    clearNodesStatusAndOutput();
  }, [clearNodesStatusAndOutput]);

  const editorControls: Control[] = useMemo(
    () => [
      {
        label: "Undo",
        icon: <IconUndo2 />,
        onClick: history.undo,
      },
      {
        label: "Redo",
        icon: <IconRedo2 />,
        onClick: history.redo,
      },
    ],
    [],
  );

  // Track when we should run or schedule after save completes
  const [shouldRunAfterSave, setShouldRunAfterSave] = useState(false);
  const [shouldScheduleAfterSave, setShouldScheduleAfterSave] = useState(false);

  // Effect to trigger runOrOpenInput or openRunInputDialog after saving completes
  useEffect(() => {
    if (!isSaving && shouldRunAfterSave) {
      runnerUIRef.current?.runOrOpenInput();
      setShouldRunAfterSave(false);
    }
    if (!isSaving && shouldScheduleAfterSave) {
      runnerUIRef.current?.openRunInputDialog();
      setShouldScheduleAfterSave(false);
    }
  }, [isSaving, shouldRunAfterSave, shouldScheduleAfterSave]);

  const handleRunButton = useCallback(async () => {
    if (isRunning) return;
    if (!savedAgent) {
      toast({
        title: `Please save the agent first, using the button in the left sidebar.`,
      });
      return;
    }
    await saveAgent();
    setShouldRunAfterSave(true);
  }, [isRunning, savedAgent, toast, saveAgent]);

  const handleScheduleButton = useCallback(async () => {
    if (isScheduling) return;
    if (!savedAgent) {
      toast({
        title: `Please save the agent first, using the button in the left sidebar.`,
      });
      return;
    }
    await saveAgent();
    setShouldScheduleAfterSave(true);
  }, [isScheduling, savedAgent, toast, saveAgent]);

  const isNewBlockEnabled = useGetFlag(Flag.NEW_BLOCK_MENU);
  const isGraphSearchEnabled = useGetFlag(Flag.GRAPH_SEARCH);

  const onDragOver = useCallback((event: React.DragEvent) => {
    event.preventDefault();
    event.dataTransfer.dropEffect = "copy";
  }, []);

  const onDrop = useCallback(
    async (event: React.DragEvent) => {
      event.preventDefault();

      const blockData = event.dataTransfer.getData("application/reactflow");
      if (!blockData) return;

      try {
        const { blockId, blockName, hardcodedValues } = JSON.parse(blockData);

        // Convert screen coordinates to flow coordinates
        const position = screenToFlowPosition({
          x: event.clientX,
          y: event.clientY,
        });

        await createAndAddNode(
          blockId,
          blockName,
          hardcodedValues || {},
          position,
        );
      } catch (error) {
        console.error("Failed to drop block:", error);
      }
    },
    [screenToFlowPosition, createAndAddNode],
  );

  const buildContextValue: BuilderContextType = useMemo(
    () => ({
      libraryAgent,
      visualizeBeads,
      setIsAnyModalOpen,
      getNextNodeId,
      getNodeTitle,
      availableFlows,
      resolutionMode,
      enterResolutionMode,
      exitResolutionMode,
      applyPendingUpdate,
    }),
    [
      libraryAgent,
      visualizeBeads,
      getNextNodeId,
      getNodeTitle,
      availableFlows,
      resolutionMode,
      enterResolutionMode,
      applyPendingUpdate,
      exitResolutionMode,
    ],
  );

  return (
    <BuilderContext.Provider value={buildContextValue}>
      <div className={className}>
        <ReactFlow
          nodes={nodes.map((node) =>
            node.id in immediateNodePositions
              ? {
                  ...node,
                  position: immediateNodePositions[node.id] || node.position,
                }
              : node,
          )}
          edges={edges}
          nodeTypes={{ custom: CustomNode }}
          edgeTypes={{ custom: CustomEdge }}
          connectionLineComponent={ConnectionLine}
          onConnect={onConnect}
          onNodesChange={onNodesChange}
          onNodesDelete={onNodesDelete}
          onEdgesChange={onEdgesChange}
          onNodeDragStop={onNodeDragEnd}
          onNodeDragStart={onNodeDragStart}
          onDrop={onDrop}
          onDragOver={onDragOver}
          deleteKeyCode={["Backspace", "Delete"]}
          minZoom={0.1}
          maxZoom={2}
          className="dark:bg-slate-900"
        >
          <Controls />
          <Background className="dark:bg-slate-800" />
          {savedAgent && (
            <FloatingSafeModeToggle
              graph={savedAgent}
              className="right-2 top-32 p-2"
            />
          )}
          {isNewBlockEnabled ? (
            <NewControlPanel
              flowExecutionID={flowExecutionID}
              visualizeBeads={visualizeBeads}
              pinSavePopover={pinSavePopover}
              pinBlocksPopover={pinBlocksPopover}
              // nodes={nodes}
              onNodeSelect={navigateToNode}
              onNodeHover={highlightNode}
            />
          ) : (
            <ControlPanel
              className="absolute z-20"
              controls={editorControls}
              topChildren={
                <>
                  <BlocksControl
                    pinBlocksPopover={pinBlocksPopover} // Pass the state to BlocksControl
                    blocks={availableBlocks}
                    addBlock={addNode}
                    flows={availableFlows}
                    nodes={nodes}
                  />
                  {isGraphSearchEnabled && (
                    <GraphSearchControl
                      nodes={nodes}
                      onNodeSelect={navigateToNode}
                      onNodeHover={highlightNode}
                    />
                  )}
                </>
              }
              botChildren={
                <SaveControl
                  agentMeta={savedAgent}
                  canSave={!isSaving && !isRunning && !isStopping}
                  onSave={saveAgent}
                  agentDescription={agentDescription}
                  onDescriptionChange={setAgentDescription}
                  agentName={agentName}
                  onNameChange={setAgentName}
                  agentRecommendedScheduleCron={agentRecommendedScheduleCron}
                  onRecommendedScheduleCronChange={
                    setAgentRecommendedScheduleCron
                  }
                  pinSavePopover={pinSavePopover}
                />
              }
            />
          )}

          {!graphHasWebhookNodes ? (
            <BuildActionBar
              className="absolute bottom-0 left-1/2 z-20 -translate-x-1/2"
              onClickAgentOutputs={runnerUIRef.current?.openRunnerOutput}
              onClickRunAgent={handleRunButton}
              onClickStopRun={stopRun}
              onClickScheduleButton={handleScheduleButton}
              isDisabled={!savedAgent}
              isRunning={isRunning}
              resolutionModeActive={resolutionMode.active}
            />
          ) : (
            <Alert className="absolute bottom-4 left-1/2 z-20 w-auto -translate-x-1/2 select-none">
              <AlertTitle>You are building a Trigger Agent</AlertTitle>
              <AlertDescription>
                Your agent{" "}
                {savedAgent?.nodes.some((node) => node.webhook)
                  ? "is listening"
                  : "will listen"}{" "}
                for its trigger and will run when the time is right.
                <br />
                You can view its activity in your{" "}
                <Link
                  href={
                    libraryAgent
                      ? `/library/agents/${libraryAgent.id}`
                      : "/library"
                  }
                  className="underline"
                >
                  Agent Library
                </Link>
                .
              </AlertDescription>
            </Alert>
          )}
        </ReactFlow>
      </div>
      {savedAgent && (
        <RunnerUIWrapper
          ref={runnerUIRef}
          graph={savedAgent}
          nodes={nodes}
          graphExecutionError={graphExecutionError}
          createRunSchedule={createRunSchedule}
          saveAndRun={saveAndRun}
        />
      )}
      <FloatingReviewsPanel
        executionId={flowExecutionID || undefined}
        graphId={flowID || undefined}
        className="fixed bottom-24 right-4"
      />
      <Suspense fallback={null}>
        <OttoChatWidget
          graphID={flowID}
          className="fixed bottom-4 right-4 z-20"
        />
      </Suspense>
    </BuilderContext.Provider>
  );
};

const WrappedFlowEditor: typeof FlowEditor = (props) => (
  <ReactFlowProvider>
    <FlowEditor {...props} />
  </ReactFlowProvider>
);

export default WrappedFlowEditor;
