Agent/frontend/react-template/src/components/FlowChart/FlowChart.tsx

import { useCallback, useEffect, useMemo, useRef, useState, forwardRef, useImperativeHandle } from "react";
import type { ForwardRefRenderFunction } from "react";
import type { Goal } from "../../types/goal";
import type { Edge as EdgeType, Message, MessageContent } from "../../types/message";
import type { AgentMode } from "../../types/trace";
import {
  visibleSubTraceEntries,
  type SubTraceEntry,
} from "./subTraceEntries";
import { goalStatusPresentation } from "./goalStatus";
import { ArrowMarkers } from "./components/ArrowMarkers";
import styles from "./styles/FlowChart.module.css";
import { ImagePreviewModal } from "../ImagePreview/ImagePreviewModal";
import { extractImagesFromMessage } from "../../utils/imageExtraction";

interface FlowChartProps {
  goals: Goal[];
  msgGroups?: Record<string, Message[]>;
  invalidBranches?: Message[][];
  onNodeClick?: (node: Goal | Message, edge?: EdgeType) => void;
  onSubTraceClick?: (parentGoal: Goal, entry: SubTraceEntry) => void;
  agentMode?: AgentMode;
}

export interface FlowChartRef {
  expandAll: () => void;
  collapseAll: () => void;
  scrollToSequence: (sequence: number) => void;
}

type FlowMessage = Message & { _injectedGoalId?: string };
type NodeMarker = "agent_call" | "agent_return" | "agent_exit" | "agent_resume";
type FoldedNodeData = { goalId: string; count: number };

const AGENT_TOOLS = new Set(["agent", "evaluate"]);

const getContentParts = (content: Message["content"]): MessageContent[] => {
  if (Array.isArray(content)) return content;
  return content && typeof content === "object" ? [content] : [];
};

const getMessageText = (message: Message): string => {
  if (typeof message.content === "string") return message.content;
  return getContentParts(message.content).map((part) => part.text || "").join("");
};

const getMessageMarker = (message: Message, previous?: NodeMarker): NodeMarker | undefined => {
  const parts = getContentParts(message.content);
  const text = getMessageText(message);
  const isTerminate = text.includes("TERMINATE");

  if (previous === "agent_exit" && !isTerminate) return "agent_resume";

  const hasAgentCall = parts.some((part) =>
    part.tool_calls?.some((call) =>
      AGENT_TOOLS.has(call.function?.name || call.name || ""),
    ),
  );
  if (hasAgentCall) return "agent_call";

  const hasAgentReturn = parts.some((part) =>
    !!part.tool_name && AGENT_TOOLS.has(part.tool_name),
  );
  if (hasAgentReturn) return "agent_return";

  const hasStructuredContent = !!message.content && typeof message.content === "object";
  const hasToolCalls = parts.some((part) => (part.tool_calls?.length || 0) > 0);
  if (message.role === "assistant" && hasStructuredContent && text && !isTerminate && !hasToolCalls) {
    return "agent_exit";
  }

  return undefined;
};

const isContinueMessage = (message: Message): boolean =>
  getContentParts(message.content).some((part) =>
    part.tool_calls?.some((call) => {
      const args = call.function?.arguments || call.arguments;
      return typeof args === "string" && args.includes("continue_from");
    }),
  );

interface LayoutNodeBase {
  id: string;
  x: number;
  y: number;
  level: number;
  parentId?: string;
  isInvalid?: boolean;
  marker?: NodeMarker;
}

type LayoutNode =
  | (LayoutNodeBase & { type: "goal"; data: Goal })
  | (LayoutNodeBase & { type: "message"; data: FlowMessage })
  | (LayoutNodeBase & { type: "folded"; data: FoldedNodeData });

interface LayoutEdge {
  id: string;
  source: LayoutNode;
  target: LayoutNode;
  type: "line" | "bracket" | "sibling_line" | "fold_link" | "arc" | "compression_link";
  level: number;
  collapsible: boolean;
  collapsed: boolean;
  children?: LayoutNode[];
  isInvalid?: boolean;
  isFoldedEdge?: boolean;
}

const FlowChartComponent: ForwardRefRenderFunction<FlowChartRef, FlowChartProps> = (
  {
    goals,
    msgGroups = {},
    invalidBranches,
    onNodeClick,
    onSubTraceClick,
    agentMode = "legacy",
  },
  ref,
) => {
  const displayGoals = useMemo(() => {
    if (!goals) return [];
    return goals.filter((g) => !g.parent_id);
  }, [goals]);

  const containerRef = useRef<HTMLDivElement>(null);
  const scrollContainerRef = useRef<HTMLDivElement>(null);
  const svgRef = useRef<SVGSVGElement>(null);

  const [dimensions, setDimensions] = useState({ width: 1200, height: 800 });
  const [selectedNodeId, setSelectedNodeId] = useState<string | null>(null);

  const [panOffset, setPanOffset] = useState({ x: 0, y: 0 });
  const [isPanning, setIsPanning] = useState(false);
  const [zoom, setZoom] = useState(1);
  const panStartRef = useRef({ x: 0, y: 0, originX: 0, originY: 0 });

  const [hoveredArcId, setHoveredArcId] = useState<string | null>(null);
  const [hoverPos, setHoverPos] = useState<{ x: number, y: number } | null>(null);
  const zoomRange = { min: 0.2, max: 2.4 };

  const [viewMode, setViewMode] = useState<"scroll" | "panzoom">("scroll");
  const [previewImage, setPreviewImage] = useState<string | null>(null);

  const clampZoom = (value: number) => Math.min(zoomRange.max, Math.max(zoomRange.min, value));
  const resetView = () => {
    setZoom(1);
    setPanOffset({ x: 0, y: 0 });
  };

  useEffect(() => {
    const element = containerRef.current;
    if (!element) return;
    const update = () => {
      const rect = element.getBoundingClientRect();
      const width = Math.max(800, rect.width || 0);
      const height = Math.max(600, rect.height || 0);
      setDimensions({ width, height });
    };
    update();
    if (typeof ResizeObserver === "undefined") {
      window.addEventListener("resize", update);
      return () => window.removeEventListener("resize", update);
    }
    const observer = new ResizeObserver(() => update());
    observer.observe(element);
    return () => observer.disconnect();
  }, []);

  useEffect(() => {
    if (!isPanning) return;
    const handleMove = (event: MouseEvent) => {
      const { x, y, originX, originY } = panStartRef.current;
      const nextX = originX + (event.clientX - x);
      const nextY = originY + (event.clientY - y);
      setPanOffset({ x: nextX, y: nextY });
    };
    const handleUp = () => setIsPanning(false);
    window.addEventListener("mousemove", handleMove);
    window.addEventListener("mouseup", handleUp);
    return () => {
      window.removeEventListener("mousemove", handleMove);
      window.removeEventListener("mouseup", handleUp);
    };
  }, [isPanning]);

  const [collapsedGoals, setCollapsedGoals] = useState<Set<string>>(new Set());

  const toggleCollapse = useCallback((goalId: string, e: React.MouseEvent) => {
    e.stopPropagation();
    setCollapsedGoals((prev) => {
      const next = new Set(prev);
      if (next.has(goalId)) next.delete(goalId);
      else next.add(goalId);
      return next;
    });
  }, []);

  const layoutData = useMemo(() => {
    if (!displayGoals || displayGoals.length === 0) return { nodes: [], edges: [] };

    const nodes: LayoutNode[] = [];
    const edges: LayoutEdge[] = [];

    const goalLevels = new Map<string, number>();
    const goalParents = new Map<string, string>();
    let maxGoalLevel = -1;

    const computeGoalHierarchy = (list: Goal[], level: number, parentId?: string) => {
      list.forEach((goal) => {
        goalLevels.set(goal.id, level);
        maxGoalLevel = Math.max(maxGoalLevel, level);
        if (parentId) {
          goalParents.set(goal.id, parentId);
        }
        if (goal.sub_goals && goal.sub_goals.length > 0) {
          computeGoalHierarchy(goal.sub_goals, level + 1, goal.id);
        }
      });
    };
    computeGoalHierarchy(displayGoals, 0);

    const isGoalCollapsed = (goalId?: string) => {
      if (!goalId) return false;
      let curr: string | undefined = goalId;
      while (curr) {
        if (collapsedGoals.has(curr)) return true;
        curr = goalParents.get(curr);
      }
      return false;
    };

    let allMessages: FlowMessage[] = [];
    Object.entries(msgGroups).forEach(([goalId, msgs]) => {
      msgs.forEach((m) => {
        allMessages.push({ ...m, id: m.id || m.message_id, _injectedGoalId: goalId });
      });
    });

    allMessages.sort((a, b) => {
      const seqA = typeof a.sequence === "number" ? a.sequence : 0;
      const seqB = typeof b.sequence === "number" ? b.sequence : 0;
      return seqA - seqB;
    });

    allMessages = allMessages.filter((msg, idx, self) => {
      const msgId = msg.id || `msg-${idx}`;
      return idx === self.findIndex((t) => (t.id || `msg-${idx}`) === msgId);
    });

    const MESSAGE_X = 200;
    const MESSAGE_Y_START = 100;
    const MESSAGE_Y_STEP = 100;
    const GOAL_X_START = 460;
    const GOAL_X_STEP = 200;

    const msgNodeMap = new Map<string, LayoutNode>();
    const foldedNodeMap = new Map<string, LayoutNode>();
    const timelineNodes: LayoutNode[] = [];

    const timelineItems: (FlowMessage | { _isFoldedNode: true; goalId: string; count: number })[] = [];
    const foldedGoalsContent = new Map<string, number>();

    let currentFold: string | null = null;

    allMessages.forEach(msg => {
      const goalId = msg._injectedGoalId || msg.goal_id;

      let highestCollapsed: string | null = null;
      let curr: string | undefined = goalId;
      while (curr) {
        if (collapsedGoals.has(curr)) highestCollapsed = curr;
        curr = goalParents.get(curr);
      }

      if (highestCollapsed) {
        foldedGoalsContent.set(highestCollapsed, (foldedGoalsContent.get(highestCollapsed) || 0) + 1);
        if (currentFold !== highestCollapsed) {
          timelineItems.push({ _isFoldedNode: true, goalId: highestCollapsed, count: 1 });
          currentFold = highestCollapsed;
        } else {
          const lastItem = timelineItems[timelineItems.length - 1] as { _isFoldedNode: true, count: number };
          if (lastItem && lastItem._isFoldedNode) {
            lastItem.count++;
          }
        }
      } else {
        timelineItems.push(msg);
        currentFold = null;
      }
    });

    let currentY = MESSAGE_Y_START;
    let previousMarker: LayoutNode["marker"] = undefined;

    timelineItems.forEach((item, i) => {
      if ("_isFoldedNode" in item) {
        const count = item.count;
        const stableId = `folded-${item.goalId}-chunk-${i}`;
        const node: LayoutNode = {
          id: stableId,
          x: MESSAGE_X,
          y: currentY,
          data: { goalId: item.goalId, count },
          type: "folded",
          level: maxGoalLevel + 1,
          parentId: item.goalId,
        };
        nodes.push(node);
        timelineNodes.push(node);
        if (!foldedNodeMap.has(item.goalId)) {
          foldedNodeMap.set(item.goalId, node);
        }
      } else {
        const msg = item as FlowMessage;
        const marker = getMessageMarker(msg, previousMarker);

        previousMarker = marker || previousMarker;

        const stableId = msg.id || `msg-${i}`;
        const node: LayoutNode = {
          id: stableId,
          x: MESSAGE_X,
          y: currentY,
          data: msg,
          type: "message",
          level: maxGoalLevel + 1,
          parentId: msg._injectedGoalId || msg.goal_id,
          marker,
        };
        nodes.push(node);
        timelineNodes.push(node);
        msgNodeMap.set(node.id, node);
      }
      currentY += MESSAGE_Y_STEP;
    });

    for (let i = 0; i < timelineNodes.length - 1; i++) {
      const source = timelineNodes[i];
      const target = timelineNodes[i + 1];

      let shouldConnect = true;
      if (target.type === "message") {
        const pSeq = target.data.parent_sequence;
        const sourceSeq = source.type === "message" ? source.data.sequence : undefined;
        if (source.type !== "folded" && typeof pSeq === "number" && sourceSeq !== pSeq) {
          shouldConnect = false;
        }
      }

      if (shouldConnect) {
        const isFoldedEdge = source.type === "folded" || target.type === "folded";
        edges.push({
          id: `seq-${source.id}-${target.id}`,
          source,
          target,
          type: "line",
          level: 0,
          collapsible: false,
          collapsed: false,
          isFoldedEdge
        });
      }
    }

    const seqNodeMap = new Map<number, LayoutNode>();
    timelineNodes.forEach((node) => {
      if (node.type === "message") {
        const seq = node.data.sequence;
        if (typeof seq === "number") seqNodeMap.set(seq, node);
      }
    });

    timelineNodes.forEach((node, i) => {
      if (node.type === "message") {
        const pSeq = node.data.parent_sequence;
        if (typeof pSeq === "number") {
          const parentNode = seqNodeMap.get(pSeq);
          if (parentNode) {
            const pIdx = timelineNodes.indexOf(parentNode);
            if (i > pIdx + 1) { // It's a non-immediate chronological jump
              edges.push({
                id: `arc-${parentNode.id}-${node.id}`,
                source: parentNode,
                target: node,
                type: "arc",
                level: 0,
                collapsible: false,
                collapsed: false,
              });
            }
          }
        }
      }
    });

    if (timelineNodes.length > 0) {
      const firstMsg = timelineNodes[0];
      const lastMsg = timelineNodes[timelineNodes.length - 1];

      const startY = firstMsg.y - MESSAGE_Y_STEP;
      const startNode: LayoutNode = { id: "SYSTEM_START", x: MESSAGE_X, y: startY, data: { description: "START" } as Goal, type: "goal", level: 0 };
      nodes.push(startNode);
      edges.push({ id: "sys-start", source: startNode, target: firstMsg, type: "line", level: 0, collapsible: false, collapsed: false });

      const hasEndGoal = displayGoals.some((goal) => goal.id === "END");
      const endY = lastMsg.y + MESSAGE_Y_STEP;
      const endNode: LayoutNode = { id: "SYSTEM_END", x: MESSAGE_X, y: endY, data: { description: hasEndGoal ? "END" : "TERMINATED" } as Goal, type: "goal", level: 0 };
      nodes.push(endNode);
      edges.push({ id: "sys-end", source: lastMsg, target: endNode, type: "line", level: 0, collapsible: false, collapsed: false });
    }

    const allGoalNodes: LayoutNode[] = [];
    const placedGoals: { x: number; y: number; minY: number; maxY: number; hasBracket: boolean }[] = [];

    const placeGoalsBottomUp = (goalsList: Goal[], parentLevel: number): LayoutNode[] => {
      const resultNodes: LayoutNode[] = [];

      goalsList.forEach(goal => {
        if (goal.id === "START" || goal.id === "END") return;
        if (goal.parent_id && isGoalCollapsed(goal.parent_id)) return;

        const isSelfCollapsed = collapsedGoals.has(goal.id);

        const level = parentLevel + 1;
        maxGoalLevel = Math.max(maxGoalLevel, level);

        let subGoalNodes: LayoutNode[] = [];
        if (!isSelfCollapsed && goal.sub_goals && goal.sub_goals.length > 0) {
          subGoalNodes = placeGoalsBottomUp(goal.sub_goals, level);
        }

        const directMsgNodes: LayoutNode[] = [];
        if (!isSelfCollapsed) {
          timelineNodes.forEach((node) => {
            if (node.type === "message") {
              const m = node.data as Message & { _injectedGoalId: string };
              if ((m._injectedGoalId === goal.id || m.goal_id === goal.id) && m.id) {
                directMsgNodes.push(node);
              }
            }
          });
        }

        let children = [...subGoalNodes, ...directMsgNodes];
        let y = MESSAGE_Y_START;

        if (isSelfCollapsed) {
          const foldedNode = foldedNodeMap.get(goal.id);
          if (foldedNode) {
            y = foldedNode.y;
          } else {
            y = allGoalNodes.length > 0 ? Math.max(...allGoalNodes.map(n => n.y)) + MESSAGE_Y_STEP : MESSAGE_Y_START;
          }
          children = [];
        } else if (children.length > 0) {
          const minY = Math.min(...children.map(c => c.y));
          const maxY = Math.max(...children.map(c => c.y));
          y = (minY + maxY) / 2;
        } else {
          y = allGoalNodes.length > 0 ? Math.max(...allGoalNodes.map(n => n.y)) + MESSAGE_Y_STEP : MESSAGE_Y_START;
        }

        const node: LayoutNode = {
          id: goal.id,
          x: 0,
          y,
          data: goal,
          type: "goal",
          level,
          parentId: goal.parent_id,
        };

        allGoalNodes.push(node);

        if (isSelfCollapsed) {
          const myFoldedChunks = timelineNodes.filter(
            (timelineNode) => timelineNode.type === "folded" && timelineNode.data.goalId === goal.id,
          );
          myFoldedChunks.forEach((chunk, idx) => {
            edges.push({
              id: `fold-link-${goal.id}-${idx}`,
              source: node,
              target: chunk,
              type: "fold_link",
              level,
              collapsible: false,
              collapsed: false
            });
          });
        } else if (children.length > 0) {
          edges.push({
            id: `bracket-${goal.id}`,
            source: node,
            target: node,
            type: "bracket",
            level: level,
            collapsible: false,
            collapsed: false,
            children: children
          });
        }

        resultNodes.push(node);
      });

      return resultNodes;
    };

    placeGoalsBottomUp(displayGoals, -1);

    timelineNodes.forEach((node) => {
      if (node.type === "message") {
        const textStr = getMessageText(node.data);
        const isCompression = textStr.includes("摘要") || textStr.includes("压缩");
        if (isCompression) {
          const goalId = node.data._injectedGoalId || node.data.goal_id;
          if (goalId) {
            const goalNode = allGoalNodes.find(g => g.id === goalId);
            if (goalNode) {
              edges.push({
                id: `comp-link-${node.id}`,
                source: node,
                target: goalNode,
                type: "compression_link",
                level: 0,
                collapsible: false,
                collapsed: false,
              });
            }
          }
        }
      }
    });

    // Compute Sibling Connections
    const childrenByParent = new Map<string | undefined, LayoutNode[]>();
    allGoalNodes.forEach(node => {
      const pId = node.parentId;
      if (!childrenByParent.has(pId)) childrenByParent.set(pId, []);
      childrenByParent.get(pId)!.push(node);
    });

    childrenByParent.forEach(group => {
      group.sort((a, b) => a.y - b.y);
      for (let i = 0; i < group.length - 1; i++) {
        edges.push({
          id: `sibling-${group[i].id}-${group[i + 1].id}`,
          source: group[i],
          target: group[i + 1],
          type: "sibling_line",
          level: group[i].level,
          collapsible: false,
          collapsed: false
        });
      }
    });

    allGoalNodes.forEach(node => {
      const isSelfCollapsed = collapsedGoals.has(node.id);
      const effectiveLevel = isSelfCollapsed ? maxGoalLevel : node.level;
      let x = GOAL_X_START + (maxGoalLevel - effectiveLevel) * GOAL_X_STEP;

      let minY = node.y;
      let maxY = node.y;
      const myBracket = edges.find(e => e.type === "bracket" && e.source.id === node.id);
      const hasBracket = !!(myBracket && myBracket.children && myBracket.children.length > 0);

      if (hasBracket) {
        minY = Math.min(...myBracket!.children!.map(c => c.y));
        maxY = Math.max(...myBracket!.children!.map(c => c.y));
      }

      let collision = true;
      let attempts = 0;
      while (collision && attempts < 10) {
        collision = false;
        for (const placed of placedGoals) {
          if (Math.abs(placed.x - x) < 10) {
            const nodeOverlap = Math.abs(placed.y - node.y) < 70;
            const bracketOverlap = hasBracket && placed.hasBracket && !(maxY + 20 < placed.minY || minY - 20 > placed.maxY);
            if (nodeOverlap || bracketOverlap) {
              collision = true;
              x += 200;
              break;
            }
          }
        }
        attempts++;
      }

      node.x = x;
      placedGoals.push({ x, y: node.y, minY, maxY, hasBracket });
      nodes.push(node);
    });

    if (invalidBranches && invalidBranches.length > 0) {
      const validMsgMap = new Map<number, LayoutNode>();
      nodes.forEach((n) => {
        if (n.type === "message") {
          const seq = n.data.sequence;
          if (typeof seq === "number") {
            validMsgMap.set(seq, n);
          }
        }
      });

      const branchesByParent = new Map<string, Message[][]>();

      invalidBranches.forEach((branch) => {
        if (branch.length === 0) return;
        const firstMsg = branch[0];
        const pSeq = firstMsg.parent_sequence;

        if (typeof pSeq === "number") {
          const parentNode = validMsgMap.get(pSeq);
          if (parentNode) {
            if (!branchesByParent.has(parentNode.id)) branchesByParent.set(parentNode.id, []);
            branchesByParent.get(parentNode.id)!.push(branch);
          }
        }
      });

      branchesByParent.forEach((branches, parentId) => {
        const parentNode = nodes.find((n) => n.id === parentId);
        if (!parentNode) return;

        branches.forEach((branch, branchIdx) => {
          const offsetX = 220 + branchIdx * 200;
          let currentParent = parentNode;
          let branchPrevMarker = parentNode.marker;

          branch.forEach((msg, idx) => {
            const stableId = msg.id || `${parentNode.id}-branch-${branchIdx}-${idx}`;
            const nodeId = `invalid-${stableId}`;

            const marker = getMessageMarker(msg, branchPrevMarker);
            branchPrevMarker = marker || branchPrevMarker;

            const node: LayoutNode = {
              id: nodeId,
              x: parentNode.x - offsetX,
              y: parentNode.y + (idx + 1) * MESSAGE_Y_STEP,
              data: msg,
              type: "message",
              level: parentNode.level,
              parentId: parentNode.id,
              isInvalid: false,
              marker,
            };
            nodes.push(node);

            edges.push({
              id: `edge-${currentParent.id}-${node.id}`,
              source: currentParent,
              target: node,
              type: "line",
              level: 0,
              collapsible: false,
              collapsed: false,
              isInvalid: false,
            });

            currentParent = node;
          });
        });
      });
    }

    return { nodes, edges };
  }, [displayGoals, msgGroups, invalidBranches, collapsedGoals]);

  const allGoalIds = useMemo(() => {
    const ids: string[] = [];
    const traverse = (list: Goal[]) => {
      list.forEach((g) => {
        ids.push(g.id);
        if (g.sub_goals) traverse(g.sub_goals);
      });
    };
    if (displayGoals) traverse(displayGoals);
    return ids;
  }, [displayGoals]);


  const visibleData = layoutData;

  const contentSize = useMemo(() => {
    if (visibleData.nodes.length === 0) return { width: 0, height: 0, minX: 0, minY: 0 };
    let minX = Infinity;
    let maxX = -Infinity;
    let minY = Infinity;
    let maxY = -Infinity;
    visibleData.nodes.forEach((node) => {
      minX = Math.min(minX, node.x);
      maxX = Math.max(maxX, node.x);
      minY = Math.min(minY, node.y);
      maxY = Math.max(maxY, node.y);
    });

    minX -= 70;
    maxX += 70;

    const centerX = (minX + maxX) / 2;
    const baseWidth = (maxX - minX) + 400;
    const baseHeight = (maxY - minY) + 300;

    const finalWidth = Math.max(dimensions.width, baseWidth);
    const finalHeight = Math.max(dimensions.height, baseHeight);

    const startX = centerX - finalWidth / 2;
    const startY = minY - 150;

    return {
      width: finalWidth,
      height: finalHeight,
      minX: startX,
      minY: startY
    };
  }, [visibleData, dimensions]);

  useImperativeHandle(ref, () => ({
    expandAll: () => setCollapsedGoals(new Set()),
    collapseAll: () => setCollapsedGoals(new Set(allGoalIds)),
    scrollToSequence: (sequence: number) => {
      const targetNode = layoutData?.nodes.find(
        (n) => n.type === "message" && (n.data as Message).sequence === sequence
      );
      if (!targetNode) return;
      setSelectedNodeId(targetNode.id);
      onNodeClick?.(targetNode.data as Message);
      if (viewMode === "scroll" && scrollContainerRef.current) {
        const rect = scrollContainerRef.current.getBoundingClientRect();
        scrollContainerRef.current.scrollTo({
          left: (targetNode.x - contentSize.minX) - rect.width / 2 + 100,
          top: (targetNode.y - contentSize.minY) - rect.height / 2,
          behavior: "smooth",
        });
      } else if (viewMode === "panzoom" && containerRef.current) {
        const rect = containerRef.current.getBoundingClientRect();
        setPanOffset({
          x: rect.width / 2 - targetNode.x * zoom - 100 * zoom,
          y: rect.height / 2 - targetNode.y * zoom,
        });
      }
    },
  }), [allGoalIds, layoutData, contentSize, viewMode, zoom, onNodeClick]);

  const toggleView = () => {
    if (viewMode === "scroll") {
      setViewMode("panzoom");
      resetView();
    } else {
      setViewMode("scroll");
      resetView();
    }
  };

  const getBracketPath = (source: LayoutNode, children: LayoutNode[] | undefined) => {
    const sx = source.x - 70;
    const sy = Math.round(source.y);
    const dropX = sx - 40;
    const endX = dropX - 30;

    if (!children || children.length === 0) {
      return `M ${sx},${sy} L ${dropX},${sy}`;
    }

    const minY = Math.round(Math.min(...children.map(c => c.y)));
    const maxY = Math.round(Math.max(...children.map(c => c.y)));

    if (children.length === 1 || Math.abs(minY - maxY) < 1) {
      return `M ${sx},${sy} L ${dropX},${sy} M ${endX},${minY} L ${dropX},${minY}`;
    }

    return `M ${sx},${sy} L ${dropX},${sy} M ${endX},${minY} L ${dropX},${minY} L ${dropX},${maxY} L ${endX},${maxY}`;
  };

  const getLinePath = (source: LayoutNode, target: LayoutNode) => {
    if (source.isInvalid && target.isInvalid) {
      return `M ${source.x},${source.y + 25} L ${target.x},${target.y - 28}`;
    }
    if (!source.isInvalid && target.isInvalid) {
      // Diverging from main timeline
      return `M ${source.x - 40},${source.y + 25} L ${target.x},${target.y - 28}`;
    }
    if (Math.abs(source.x - target.x) < 1) {
      if (source.id === "SYSTEM_START" || target.id === "SYSTEM_END") {
        return `M ${source.x},${source.y + 25} L ${target.x},${target.y - 28}`;
      }
      return `M ${source.x},${source.y + 25} L ${target.x},${target.y - 28}`;
    } else {
      if (source.x > target.x) {
        return `M ${source.x - 70},${source.y} L ${target.x + 75},${target.y}`;
      } else {
        return `M ${source.x + 70},${source.y} L ${target.x - 75},${target.y}`;
      }
    }
  };

  const getBracketColor = (level: number) => {
    if (level === 0) return "#10b981"; // Emerald
    if (level === 1) return "#ef4444"; // Red
    if (level === 2) return "#3b82f6"; // Blue
    return "#f59e0b"; // Amber
  };

  const getStrokeWidth = (level: number, type: string) => {
    if (type === "bracket") return Math.max(1.5, 4 - level * 1);
    return 2;
  };

  const handleNodeClick = useCallback(
    (node: LayoutNode, e: React.MouseEvent) => {
      e.stopPropagation();
      setSelectedNodeId(node.id);
      if (node.type === "goal") {
        onNodeClick?.(node.data as Goal);
      } else if (node.type === "message") {
        onNodeClick?.(node.data as Message);
      }
    },
    [onNodeClick],
  );

  const jumpToNode = (nodeId: string, e: React.MouseEvent) => {
    e.stopPropagation();
    const targetNode = layoutData?.nodes.find((node) => node.id === nodeId);
    if (!targetNode) return;

    setSelectedNodeId(targetNode.id);
    if (targetNode.type === "goal") {
      onNodeClick?.(targetNode.data as Goal);
    } else if (targetNode.type === "message") {
      onNodeClick?.(targetNode.data as Message);
    }

    if (viewMode === "scroll" && scrollContainerRef.current) {
      const rect = scrollContainerRef.current.getBoundingClientRect();
      scrollContainerRef.current.scrollTo({
        left: (targetNode.x - contentSize.minX) - rect.width / 2 + 100,
        top: (targetNode.y - contentSize.minY) - rect.height / 2,
        behavior: "smooth"
      });
    } else if (viewMode === "panzoom" && containerRef.current) {
      const rect = containerRef.current.getBoundingClientRect();
      setPanOffset({
        x: rect.width / 2 - targetNode.x * zoom - 100 * zoom,
        y: rect.height / 2 - targetNode.y * zoom
      });
    }
  };

  if (!layoutData) return <div>Loading...</div>;

  return (
    <div className={styles.container} ref={containerRef}>
      <div
        ref={scrollContainerRef}
        className={styles.scrollContainer}
        style={{ overflow: viewMode === "scroll" ? "auto" : "hidden" }}
      >
        <svg
          ref={svgRef}
          viewBox={
            viewMode === "scroll"
              ? `${contentSize.minX} ${contentSize.minY} ${contentSize.width} ${contentSize.height}`
              : `0 0 ${dimensions.width} ${dimensions.height}`
          }
          preserveAspectRatio="xMidYMid meet"
          className={`${styles.svg} ${isPanning ? styles.panning : ""}`}
          style={{
            cursor: viewMode === "scroll" ? "default" : isPanning ? "grabbing" : "grab",
            width: viewMode === "scroll" ? contentSize.width : "100%",
            height: viewMode === "scroll" ? contentSize.height : "100%",
          }}
          onWheel={(event) => {
            if (viewMode === "scroll") return;
            event.preventDefault();
            const rect = svgRef.current?.getBoundingClientRect();
            if (!rect) return;
            const cursorX = event.clientX - rect.left;
            const cursorY = event.clientY - rect.top;
            const nextZoom = clampZoom(zoom * (event.deltaY > 0 ? 0.92 : 1.08));
            if (nextZoom === zoom) return;
            const scale = nextZoom / zoom;
            setPanOffset((prev) => ({
              x: cursorX - (cursorX - prev.x) * scale,
              y: cursorY - (cursorY - prev.y) * scale,
            }));
            setZoom(nextZoom);
          }}
          onDoubleClick={() => resetView()}
          onMouseDown={(event) => {
            if (viewMode === "scroll") return;
            if (event.button !== 0) return;
            const target = event.target as Element;
            if (target.closest(`.${styles.nodes}`) || target.closest(`.${styles.links}`)) return;
            setSelectedNodeId(null);
            panStartRef.current = {
              x: event.clientX,
              y: event.clientY,
              originX: panOffset.x,
              originY: panOffset.y,
            };
            setIsPanning(true);
          }}
        >
          <defs>
            <ArrowMarkers />
          </defs>

          <g transform={viewMode === "scroll" ? undefined : `translate(${panOffset.x},${panOffset.y}) scale(${zoom})`}>
            <g className={styles.links}>
              {visibleData.edges.map((edge) => {
                if (edge.type === "sibling_line") {
                  return (
                    <path
                      key={edge.id}
                      d={`M ${edge.source.x},${edge.source.y + 25} L ${edge.target.x},${edge.target.y - 25}`}
                      fill="none"
                      stroke="#cbd5e1"
                      strokeWidth={2}
                      strokeDasharray="6,6"
                      markerEnd="url(#arrow-default)"
                    />
                  );
                }

                if (edge.type === "arc") {
                  const isSelected = edge.source.id === selectedNodeId || edge.target.id === selectedNodeId;
                  if (!isSelected && hoveredArcId !== edge.id) return null;

                  const sx = edge.source.x - 70;
                  const sy = edge.source.y;
                  const tx = edge.target.x - 70;
                  const ty = edge.target.y;

                  const dist = Math.abs(ty - sy);
                  const bulge = 40 + Math.min(200, dist / 4);
                  const arcPath = `M ${sx},${sy} C ${sx - bulge},${sy} ${tx - bulge},${ty} ${tx},${ty}`;

                  return (
                    <g
                      key={edge.id}
                      onMouseEnter={(e) => {
                        if (hoveredArcId !== edge.id) {
                          setHoveredArcId(edge.id);
                          const svg = e.currentTarget.closest("svg");
                          const ctm = (e.currentTarget as SVGGElement).getScreenCTM();
                          if (svg && ctm) {
                            const point = svg.createSVGPoint();
                            point.x = e.clientX;
                            point.y = e.clientY;
                            const localPoint = point.matrixTransform(ctm.inverse());
                            setHoverPos({ x: localPoint.x, y: localPoint.y });
                          }
                        }
                      }}
                      onMouseLeave={() => {
                        setHoveredArcId(null);
                        setHoverPos(null);
                      }}
                      style={{ cursor: "pointer" }}
                    >
                      <path
                        d={arcPath}
                        fill="none"
                        stroke="#94a3b8"
                        strokeWidth={hoveredArcId === edge.id ? 4 : 3}
                        strokeDasharray={hoveredArcId === edge.id ? undefined : "5,5"}
                        markerEnd="url(#arrow-default)"
                        style={{ pointerEvents: "none" }}
                      />
                      {/* Track Mouse ONLY while on the line */}
                      <path
                        d={arcPath}
                        fill="none"
                        stroke="transparent"
                        strokeWidth={30}
                        style={{ pointerEvents: "stroke" }}
                      />
                      {hoveredArcId === edge.id && hoverPos && (
                        <g transform={`translate(${hoverPos.x - 14}, ${hoverPos.y + 5})`}>
                          <rect width={70} height={24} fill="#94a3b8" rx={4} onClick={(e) => jumpToNode(edge.source.id, e)} />
                          <text x={35} y={16} fill="white" textAnchor="middle" style={{ fontSize: 12, pointerEvents: "none" }}>向上跳转</text>

                          <rect y={30} width={70} height={24} fill="#94a3b8" rx={4} onClick={(e) => jumpToNode(edge.target.id, e)} />
                          <text x={35} y={46} fill="white" textAnchor="middle" style={{ fontSize: 12, pointerEvents: "none" }}>向下跳转</text>
                        </g>
                      )}
                    </g>
                  );
                }

                if (edge.type === "fold_link") {
                  const sx = edge.source.x - 70;
                  const sy = edge.source.y;
                  const tx = edge.target.x + 22;
                  const ty = edge.target.y;
                  const midX = sx - 30;
                  const foldPathStr = `M ${sx},${sy} L ${midX},${sy} L ${midX},${ty} L ${tx},${ty}`;

                  return (
                    <g key={edge.id} style={{ cursor: "pointer" }} onClick={(e) => toggleCollapse(edge.source.id, e)}>
                      <path d={foldPathStr} fill="none" stroke="transparent" strokeWidth={24} />
                      <path
                        d={foldPathStr}
                        fill="none"
                        stroke="#10b981"
                        strokeWidth={2}
                        strokeDasharray="4,4"
                      />
                    </g>
                  );
                }

                if (edge.type === "compression_link") {
                  // horizontal, then vertical, then horizontal to goal
                  const sx = edge.source.x + 70;
                  const sy = edge.source.y;
                  const tx = edge.target.x - 70;
                  const ty = edge.target.y;

                  const midX = sx + 40;
                  const pathStr = `M ${sx},${sy} L ${midX},${sy} L ${midX},${ty} L ${tx},${ty}`;

                  return (
                    <path
                      key={edge.id}
                      d={pathStr}
                      fill="none"
                      stroke="#94a3b8"
                      strokeWidth={1.5}
                      strokeDasharray="4,4"
                      markerEnd="url(#arrow-default)"
                    />
                  );
                }

                const isBracket = edge.type === "bracket";
                let pathStr = "";
                let strokeWidth = 2;
                let color = "#94a3b8";

                if (isBracket) {
                  pathStr = getBracketPath(edge.source, edge.children);
                  strokeWidth = getStrokeWidth(edge.level, edge.type);
                  color = getBracketColor(edge.level);
                } else if (edge.isFoldedEdge) {
                  pathStr = getLinePath(edge.source, edge.target);
                  strokeWidth = 3;
                  color = "#10b981";
                } else {
                  pathStr = getLinePath(edge.source, edge.target);
                }

                if (isBracket) {
                  return (
                    <g key={edge.id} style={{ cursor: "pointer" }} onClick={(e) => toggleCollapse(edge.source.id, e)}>
                      <path d={pathStr} fill="none" stroke="transparent" strokeWidth={24} />
                      <path
                        d={pathStr}
                        fill="none"
                        stroke={edge.isInvalid ? "#e2e8f0" : color}
                        strokeWidth={strokeWidth}
                        strokeDasharray={edge.isInvalid ? "5,5" : undefined}
                        strokeLinejoin="round"
                      />
                    </g>
                  );
                }

                return (
                  <path
                    key={edge.id}
                    d={pathStr}
                    fill="none"
                    stroke={edge.isInvalid ? "#e2e8f0" : color}
                    strokeWidth={strokeWidth}
                    strokeDasharray={edge.isInvalid ? "5,5" : undefined}
                    markerEnd={edge.isInvalid || edge.isFoldedEdge ? undefined : "url(#arrow-default)"}
                    strokeLinejoin="round"
                    style={edge.isFoldedEdge ? { cursor: "pointer" } : undefined}
                    onClick={edge.isFoldedEdge ? (e) => toggleCollapse(edge.source.parentId || "", e) : undefined}
                  />
                );
              })}
            </g>

            <g className={styles.nodes}>
              {visibleData.nodes.map((node) => {
                if (node.type === "folded") {
                  const count = node.data.count;
                  const goalId = node.data.goalId;
                  return (
                    <g
                      key={node.id}
                      transform={`translate(${node.x},${node.y})`}
                      style={{ cursor: "pointer" }}
                      onClick={(e) => toggleCollapse(goalId, e)}
                    >
                      <circle r={14} cx={0} cy={0} fill="#ecfdf5" stroke="#10b981" strokeWidth={2} />
                      <text x={0} y={4} fontSize={11} fontWeight="bold" fill="#047857" textAnchor="middle">
                        {count}
                      </text>
                    </g>
                  );
                }

                const isGoal = node.type === "goal";
                const data = node.data as Goal;
                const subTraceEntries = isGoal
                  ? visibleSubTraceEntries(data, agentMode)
                  : [];
                const statusPresentation = isGoal
                  ? goalStatusPresentation(data.status)
                  : undefined;
                let text = isGoal ? data.description : (node.data as Message).description || "";

                let thumbnail: string | null = null;
                if (node.type === "message") {
                  const images = extractImagesFromMessage(node.data as Message);
                  if (images.length > 0) thumbnail = images[0].url;
                }

                let textColor = "#3b82f6";
                if (node.type === "message") {
                  textColor = "#64748b";
                  if (node.marker === "agent_call") textColor = "#f59e0b";
                  else if (node.marker === "agent_return") textColor = "#8b5cf6";
                  else if (node.marker === "agent_exit") textColor = "#ef4444";
                  else if (node.marker === "agent_resume") textColor = "#06b6d4";
                } else if (node.type === "goal" && node.level === 0) {
                  textColor = "#10b981";
                } else if (node.type === "goal" && node.level === 1) {
                  textColor = "#ef4444";
                }

                if (node.isInvalid) textColor = "#94a3b8";

                let emoji = "";
                if (node.type === "message" && node.id !== "SYSTEM_START" && node.id !== "SYSTEM_END") {
                  if (isContinueMessage(node.data)) {
                    emoji = "🔄🤖 ";
                    text = "Agent 重调";
                  } else if (node.marker === "agent_call") {
                    emoji = "🤖 ";
                    text = "Agent 调用";
                  } else if (node.marker === "agent_return") {
                    emoji = "✅ ";
                    text = "Agent 返回";
                  } else if (node.marker === "agent_exit") {
                    emoji = "🏁 ";
                    text = "Agent 结束";
                  } else if (node.marker === "agent_resume") {
                    emoji = "🔄 ";
                    text = "Agent 续跑: " + text;
                  }
                }

                let nodeFill = isGoal ? "#eff6ff" : "#f8fafc";
                let nodeStroke = selectedNodeId === node.id ? "#3b82f6" : node.isInvalid ? "#cbd5e1" : "#e2e8f0";

                if (isGoal) {
                  if (node.level === 0) { nodeFill = "#ecfdf5"; nodeStroke = "#10b981"; }
                  else if (node.level === 1) { nodeFill = "#fef2f2"; nodeStroke = "#ef4444"; }
                  else if (node.level === 2) { nodeFill = "#eff6ff"; nodeStroke = "#3b82f6"; }
                  if (statusPresentation) {
                    nodeFill = statusPresentation.fill;
                    nodeStroke = statusPresentation.stroke;
                    textColor = statusPresentation.text;
                  }
                }

                if (node.id === "SYSTEM_START" || node.id === "SYSTEM_END") {
                  nodeFill = "#f8fafc"; nodeStroke = "#cbd5e1"; textColor = "#64748b";
                }

                if (node.marker === "agent_call") {
                  nodeFill = "#fffbeb";
                  nodeStroke = selectedNodeId === node.id ? "#f59e0b" : "#fcd34d";
                } else if (node.marker === "agent_return") {
                  nodeFill = "#f5f3ff";
                  nodeStroke = selectedNodeId === node.id ? "#8b5cf6" : "#c4b5fd";
                } else if (node.marker === "agent_exit") {
                  nodeFill = "#fef2f2";
                  nodeStroke = selectedNodeId === node.id ? "#ef4444" : "#fca5a5";
                } else if (node.marker === "agent_resume") {
                  nodeFill = "#ecfeff";
                  nodeStroke = selectedNodeId === node.id ? "#06b6d4" : "#67e8f9";
                }

                // Detect remote children (targeted by arcs)
                const hasRemoteParent = visibleData.edges.some((e) => e.type === "arc" && e.target.id === node.id);

                return (
                  <g key={node.id} transform={`translate(${node.x},${node.y})`} onClick={(e) => handleNodeClick(node, e)}>
                    {hasRemoteParent && node.type === "message" && (
                      <rect
                        x={-73} y={-28} width={146} height={56} rx={10}
                        fill="none" stroke="#94a3b8" strokeWidth={2} strokeDasharray="4,2"
                        style={{ pointerEvents: "none" }}
                      />
                    )}
                    <rect
                      x={-70} y={-25} width={140} height={50} rx={8}
                      fill={nodeFill} stroke={selectedNodeId === node.id ? "#8b5cf6" : nodeStroke} strokeWidth={selectedNodeId === node.id ? 2 : 1}
                      strokeDasharray={node.isInvalid ? "5,5" : undefined}
                      style={{
                        filter: selectedNodeId === node.id ? "drop-shadow(0 4px 6px rgb(59 130 246 / 0.3))" : "drop-shadow(0 1px 2px rgb(0 0 0 / 0.05))"
                      }}
                    />
                    <foreignObject x={-70} y={-25} width={140} height={50}>
                      <div className="w-full h-full overflow-hidden flex items-center justify-between px-2 gap-2" style={{ color: textColor }}>
                        {isGoal && node.id !== "SYSTEM_START" && node.id !== "SYSTEM_END" && (
                          <div
                            onClick={(e) => toggleCollapse(node.id, e)}
                            className="text-[10px] w-4 h-4 flex items-center justify-center border rounded bg-white text-gray-500 hover:bg-gray-100 flex-shrink-0"
                            style={{ pointerEvents: "auto" }}
                          >
                            {collapsedGoals.has(node.id) ? '+' : '-'}
                          </div>
                        )}
                        <span className="text-xs font-semibold line-clamp-3 w-full text-center" style={{ justifyContent: thumbnail ? "flex-start" : "center", textAlign: thumbnail ? "left" : "center" }}>{emoji}{text}</span>
                        {thumbnail && (
                          <img src={thumbnail} alt="thumb" className="w-8 h-8 object-cover rounded border border-gray-200 bg-white flex-shrink-0 hover:scale-110 transition-transform" onClick={(e) => { e.stopPropagation(); setPreviewImage(thumbnail as string); }} />
                        )}
                      </div>
                    </foreignObject>
                    {node.type === "message" && (
                      <g transform="translate(-70, -25)">
                        <circle r={10} cx={0} cy={0} fill="#f1f5f9" stroke={nodeStroke} strokeWidth={1} />
                        <text x={0} y={3} fontSize={10} fontWeight="bold" fill={textColor} textAnchor="middle">
                          {node.data.sequence ?? ""}
                        </text>
                      </g>
                    )}
                    {statusPresentation && (
                      <g transform="translate(0, -35)" style={{ pointerEvents: "none" }}>
                        <rect
                          x={-42}
                          y={-8}
                          width={84}
                          height={16}
                          rx={8}
                          fill={statusPresentation.fill}
                          stroke={statusPresentation.stroke}
                        />
                        <text
                          x={0}
                          y={4}
                          fontSize={9}
                          fontWeight="bold"
                          fill={statusPresentation.text}
                          textAnchor="middle"
                        >
                          {statusPresentation.label}
                        </text>
                      </g>
                    )}
                    {subTraceEntries.map((entry, index) => {
                      const spacing = 22;
                      const startX = -((subTraceEntries.length - 1) * spacing) / 2;
                      return (
                        <g
                          key={entry.id}
                          transform={`translate(${startX + index * spacing}, 36)`}
                          style={{ cursor: "pointer" }}
                          onClick={(event) => {
                            event.stopPropagation();
                            onSubTraceClick?.(data, entry);
                          }}
                        >
                          <title>{entry.mission || entry.id}</title>
                          <rect
                            x={-10}
                            y={-8}
                            width={20}
                            height={16}
                            rx={5}
                            fill="#fffbeb"
                            stroke="#f59e0b"
                            strokeWidth={1}
                          />
                          <text
                            x={0}
                            y={4}
                            fontSize={9}
                            fontWeight="bold"
                            fill="#b45309"
                            textAnchor="middle"
                          >
                            A{index + 1}
                          </text>
                        </g>
                      );
                    })}
                  </g>
                );
              })}
            </g>
          </g>
        </svg>
      </div>

      <div className={styles.controls}>
        {viewMode === "panzoom" && (
          <>
            <button type="button" className={styles.controlButton} onClick={() => setZoom((prev) => clampZoom(prev * 1.1))}>+</button>
            <button type="button" className={styles.controlButton} onClick={() => setZoom((prev) => clampZoom(prev * 0.9))}>−</button>
            <button type="button" className={styles.controlButton} onClick={resetView}>复位</button>
          </>
        )}
        <button type="button" className={styles.controlButton} onClick={toggleView}>
          {viewMode === "scroll" ? "切换拖拽" : "切换滚动"}
        </button>
      </div>

      <ImagePreviewModal visible={!!previewImage} onClose={() => setPreviewImage(null)} src={previewImage || ""} />
    </div>
  );
};

export const FlowChart = forwardRef(FlowChartComponent);