src/components/PipelineGraph.tsx

"use client";

import { useEffect, useMemo, useState } from "react";
import {
  ReactFlow,
  Background,
  BackgroundVariant,
  Controls,
  type Edge,
  type Node,
} from "@xyflow/react";
import "@xyflow/react/dist/style.css";
import type { ImplementationMode, Pipeline } from "@/lib/types";
import {
  computeLayout,
  computeBodyGroupLayout,
  BODY_GROUP_LABEL,
  BODY_GROUP_COLOR,
  type LayoutMode,
} from "@/lib/graph-layout";
import StageNode from "./nodes/StageNode";
import GroupLabelNode from "./nodes/GroupLabelNode";

const nodeTypes = { stage: StageNode, "group-label": GroupLabelNode };
const LANDING_STEP_MS = 70;
const GROUP_TRANSITION_MS = 950;

export default function PipelineGraph({
  pipeline,
  expandedIds,
  selectedId,
  activeModes,
  showMetrics,
  layoutMode,
  groupedByLayer,
  compact = false,
  onSelect,
  onToggleExpand,
  onSetMode,
  onDeepDive,
}: {
  pipeline: Pipeline;
  expandedIds: Set<string>;
  selectedId: string | null;
  activeModes: Record<string, ImplementationMode>;
  showMetrics: boolean;
  layoutMode: LayoutMode;
  groupedByLayer: boolean;
  compact?: boolean;
  onSelect: (id: string) => void;
  onToggleExpand: (id: string) => void;
  onSetMode: (id: string, mode: ImplementationMode) => void;
  onDeepDive: (id: string) => void;
}) {
  const [transitioning, setTransitioning] = useState(false);

  // Smooth glide, not a remount: when grouping toggles, briefly enable a CSS
  // transition on node transforms, then drop it so normal dragging stays crisp.
  useEffect(() => {
    setTransitioning(true);
    const t = setTimeout(() => setTransitioning(false), GROUP_TRANSITION_MS);
    return () => clearTimeout(t);
  }, [groupedByLayer]);

  const { nodes, edges } = useMemo(() => {
    const nodes: Node[] = [];
    let landingIndex = 0;

    if (groupedByLayer) {
      const { positions, groupHeaders } = computeBodyGroupLayout(pipeline);

      groupHeaders.forEach((h) => {
        nodes.push({
          id: `group-header-${h.group}`,
          type: "group-label",
          position: { x: h.x - 160, y: h.y },
          data: { label: BODY_GROUP_LABEL[h.group], color: BODY_GROUP_COLOR[h.group] },
          draggable: false,
          selectable: false,
        });
      });

      pipeline.nodes.forEach((n) => {
        nodes.push({
          id: n.id,
          type: "stage",
          position: positions[n.id] ?? { x: 0, y: 0 },
          data: {
            node: n,
            isChild: false,
            isHub: false,
            expanded: false,
            hasChildren: false,
            showMetrics,
            activeMode: activeModes[n.id] ?? n.implementations[0]?.mode ?? "local",
            landingDelayMs: landingIndex++ * LANDING_STEP_MS,
            onSelect,
            onSetMode,
            onDeepDive,
            selected: selectedId === n.id,
          },
          draggable: true,
        });
      });

      const edges: Edge[] = pipeline.edges
        .filter((e) => e.kind !== "hosts")
        .map((e) => ({
          id: `e-${e.source}-${e.target}`,
          source: e.source,
          target: e.target,
          label: e.label,
          animated: e.kind === "data",
          type: "smoothstep",
          style:
            e.kind === "loop"
              ? { stroke: "#2dd4bf", strokeDasharray: "2 4", strokeWidth: 1.5 }
              : { stroke: "#f0b429", strokeWidth: 2, filter: "drop-shadow(0 0 3px rgba(240,180,41,0.55))" },
          labelStyle: { fill: "#e6c878", fontSize: 10, fontWeight: 600 },
          labelBgStyle: { fill: "#0e1826", fillOpacity: 0.85 },
        }));

      return { nodes, edges };
    }

    const { positions, flatChildren } = computeLayout(pipeline, expandedIds, layoutMode, compact);

    pipeline.nodes.forEach((n) => {
      nodes.push({
        id: n.id,
        type: "stage",
        position: positions[n.id] ?? { x: 0, y: 0 },
        data: {
          node: n,
          isChild: false,
          isHub: n.category === "orchestration",
          expanded: expandedIds.has(n.id),
          hasChildren: !!n.children?.length,
          showMetrics,
          compact,
          activeMode: activeModes[n.id] ?? n.implementations[0]?.mode ?? "local",
          landingDelayMs: landingIndex++ * LANDING_STEP_MS,
          onToggleExpand,
          onSelect,
          onSetMode,
          onDeepDive,
          selected: selectedId === n.id,
        },
        draggable: true,
      });
    });

    Object.entries(flatChildren).forEach(([id, child]) => {
      nodes.push({
        id,
        type: "stage",
        position: positions[id] ?? { x: 0, y: 0 },
        data: {
          node: child,
          isChild: true,
          isHub: false,
          expanded: false,
          hasChildren: false,
          showMetrics,
          activeMode: activeModes[id] ?? child.implementations[0]?.mode ?? "local",
          landingDelayMs: landingIndex++ * LANDING_STEP_MS,
          onSelect,
          onSetMode,
          onDeepDive,
          selected: selectedId === id,
        },
        draggable: true,
      });
    });

    // Grid mode is for inspecting nodes individually, not reading the flow —
    // skip edges entirely rather than let handoff labels clutter a grid.
    const edges: Edge[] = layoutMode === "grid"
      ? []
      : pipeline.edges.map((e) => ({
          id: `e-${e.source}-${e.target}`,
          source: e.source,
          target: e.target,
          label: e.label,
          animated: e.kind === "data",
          type: "smoothstep",
          style:
            e.kind === "hosts"
              ? { stroke: "#475569", strokeDasharray: "3 4", opacity: 0.45 }
              : e.kind === "loop"
              ? { stroke: "#2dd4bf", strokeDasharray: "2 4", strokeWidth: 1.5 }
              : { stroke: "#f0b429", strokeWidth: 2, filter: "drop-shadow(0 0 3px rgba(240,180,41,0.55))" },
          labelStyle: { fill: "#e6c878", fontSize: 10, fontWeight: 600 },
          labelBgStyle: { fill: "#0e1826", fillOpacity: 0.85 },
        }));

    Object.keys(flatChildren).forEach((childId) => {
      const parent = pipeline.nodes.find((n) =>
        n.children?.some((c) => c.id === childId),
      );
      if (parent) {
        edges.push({
          id: `e-${parent.id}-${childId}`,
          source: parent.id,
          target: childId,
          type: "straight",
          style: { stroke: "#64748b", strokeDasharray: "1 3", opacity: 0.6 },
        });
      }
    });

    return { nodes, edges };
  }, [pipeline, expandedIds, selectedId, activeModes, showMetrics, layoutMode, groupedByLayer, compact, onSelect, onToggleExpand, onSetMode, onDeepDive]);

  return (
    <ReactFlow
      key={layoutMode}
      nodes={nodes}
      edges={edges}
      nodeTypes={nodeTypes}
      fitView
      fitViewOptions={{ padding: 0.35 }}
      minZoom={0.3}
      maxZoom={1.5}
      proOptions={{ hideAttribution: true }}
      className={transitioning ? "grouping-transition" : undefined}
    >
      <Background variant={BackgroundVariant.Dots} gap={26} size={1} color="#1c2b3f" />
      <Controls
        showInteractive={false}
        className="!bg-[var(--panel)] !border !border-[var(--panel-border)] !fill-white [&_button]:!border-[var(--panel-border)] [&_button]:!text-white"
      />
    </ReactFlow>
  );
}