polaris-task-force/src/components/admin/narrative-flow/NarrativeFlowEditor.tsx

"use client";

import { useCallback, useEffect, useRef, useState } from "react";
import {
  ReactFlow,
  Background,
  Controls,
  MiniMap,
  useNodesState,
  useEdgesState,
  addEdge,
  type Node,
  type Edge,
  type Connection,
} from "@xyflow/react";
import "@xyflow/react/dist/style.css";

import { ContextMenu as ContextMenuPrimitive } from "radix-ui";
import { Trash2Icon } from "lucide-react";

import { useField } from "@payloadcms/ui";
import { NarrativeBeatNode } from "./NarrativeBeatNode";
import { ChoiceNode } from "./ChoiceNode";
import { OutcomeNode } from "./OutcomeNode";
import { WaitNode } from "./WaitNode";
import { NumberNode } from "./NumberNode";
import { StringNode } from "./StringNode";
import { RandomNode } from "./RandomNode";
import { BooleanNode } from "./BooleanNode";
import { DateNode } from "./DateNode";
import { GameDataNode } from "./GameDataNode";
import { FlagNode } from "./FlagNode";
import { CompareNode } from "./CompareNode";
import { MathNode } from "./MathNode";
import { FormatNode } from "./FormatNode";
import { LogicNode } from "./LogicNode";
import { NotNode } from "./NotNode";
import { BranchNode } from "./BranchNode";
import { SequenceNode } from "./SequenceNode";
import { NodeEditDialog } from "./NodeEditDialog";
import { NODE_SOCKETS, type SocketType } from "./types";

const nodeTypes = {
  narrativeBeat: NarrativeBeatNode,
  choice: ChoiceNode,
  outcome: OutcomeNode,
  wait: WaitNode,
  number: NumberNode,
  string: StringNode,
  random: RandomNode,
  boolean: BooleanNode,
  date: DateNode,
  gameData: GameDataNode,
  flag: FlagNode,
  compare: CompareNode,
  math: MathNode,
  format: FormatNode,
  logic: LogicNode,
  not: NotNode,
  branch: BranchNode,
  sequence: SequenceNode,
};

type FlowData = { nodes: Node[]; edges: Edge[] };

const defaultFlow: FlowData = { nodes: [], edges: [] };

type ContextMenuState = {
  x: number;
  y: number;
  nodeId?: string;
  edgeId?: string;
} | null;

export function NarrativeFlowEditor() {
  const { value, setValue } = useField<FlowData>();

  const initialData: FlowData =
    value && typeof value === "object" && Array.isArray((value as FlowData).nodes)
      ? (value as FlowData)
      : defaultFlow;

  const [nodes, setNodes, onNodesChange] = useNodesState(initialData.nodes);
  const [edges, setEdges, onEdgesChange] = useEdgesState(initialData.edges);

  const [contextMenu, setContextMenu] = useState<ContextMenuState>(null);
  const [editingNodeId, setEditingNodeId] = useState<string | null>(null);

  const lastSaved = useRef<string>("");

  useEffect(() => {
    const flow = { nodes, edges };
    const serialized = JSON.stringify(flow);
    if (serialized !== lastSaved.current) {
      lastSaved.current = serialized;
      setValue(flow);
    }
  }, [nodes, edges, setValue]);

  const onConnect = useCallback(
    (connection: Connection) => {
      setEdges((eds) => addEdge({ ...connection, animated: true }, eds));
    },
    [setEdges],
  );

  const isValidConnection = useCallback(
    (connection: Edge | Connection) => {
      if (!connection.source || !connection.target || !connection.sourceHandle || !connection.targetHandle) {
        return false;
      }
      const sourceNode = nodes.find((n) => n.id === connection.source);
      const targetNode = nodes.find((n) => n.id === connection.target);
      if (!sourceNode || !targetNode) return false;

      const sourceSockets = NODE_SOCKETS[sourceNode.type as string];
      const targetSockets = NODE_SOCKETS[targetNode.type as string];
      if (!sourceSockets || !targetSockets) return false;

      const sourceSocket = sourceSockets.outputs.find((s) => s.id === connection.sourceHandle);
      const targetSocket = targetSockets.inputs.find((s) => s.id === connection.targetHandle);

      // Dynamic handles (arg-*, option-*) are treated as "any"
      const sourceType = sourceSocket?.type ?? (connection.sourceHandle.startsWith("arg-") || connection.sourceHandle.startsWith("option-") ? "any" : null);
      const targetType = targetSocket?.type ?? (connection.targetHandle.startsWith("arg-") || connection.targetHandle.startsWith("option-") ? "any" : null);

      if (sourceType === null || targetType === null) return false;
      if (sourceType === "any" || targetType === "any") return true;
      return sourceType === targetType;
    },
    [nodes],
  );

  const onNodeClick = useCallback((_e: React.MouseEvent, node: Node) => {
    setEditingNodeId(node.id);
  }, []);

  const onNodeContextMenu = useCallback(
    (e: React.MouseEvent, node: Node) => {
      e.preventDefault();
      setContextMenu({ x: e.clientX, y: e.clientY, nodeId: node.id });
    },
    [],
  );

  const onEdgeContextMenu = useCallback(
    (e: React.MouseEvent, edge: Edge) => {
      e.preventDefault();
      setContextMenu({ x: e.clientX, y: e.clientY, edgeId: edge.id });
    },
    [],
  );

  const onPaneContextMenu = useCallback((e: MouseEvent | React.MouseEvent) => {
    e.preventDefault();
    setContextMenu(null);
  }, []);

  const deleteSelected = useCallback(() => {
    if (!contextMenu) return;
    if (contextMenu.nodeId) {
      setNodes((nds) => nds.filter((n) => n.id !== contextMenu.nodeId));
      setEdges((eds) =>
        eds.filter((e) => e.source !== contextMenu.nodeId && e.target !== contextMenu.nodeId),
      );
    } else if (contextMenu.edgeId) {
      setEdges((eds) => eds.filter((e) => e.id !== contextMenu.edgeId));
    }
    setContextMenu(null);
  }, [contextMenu, setNodes, setEdges]);

  const handleNodeSave = useCallback(
    (nodeId: string, data: Record<string, unknown>) => {
      setNodes((nds) =>
        nds.map((n) => (n.id === nodeId ? { ...n, data: { ...n.data, ...data } } : n)),
      );
    },
    [setNodes],
  );

  const selectedNode = editingNodeId ? nodes.find((n) => n.id === editingNodeId) ?? null : null;

  const addNode = useCallback(
    (
      type:
        | "narrativeBeat"
        | "choice"
        | "outcome"
        | "wait"
        | "number"
        | "string"
        | "random"
        | "boolean"
        | "date"
        | "gameData"
        | "flag"
        | "compare"
        | "math"
        | "format"
        | "logic"
        | "not"
        | "branch"
        | "sequence",
    ) => {
      const id = crypto.randomUUID();
      const pos = { x: Math.random() * 400, y: Math.random() * 300 };
      const defaults: Record<string, Record<string, unknown>> = {
        narrativeBeat: { title: "New beat", text: "" },
        choice: { question: "New choice?", options: [] },
        outcome: { title: "New outcome", text: "", effects: [] },
        wait: { ticks: 10 },
        number: { value: 0 },
        string: { value: "" },
        random: { min: 0, max: 100 },
        boolean: { value: false },
        date: { value: "" },
        gameData: { field: "structure_health" },
        flag: { flagName: "" },
        compare: { operator: ">" },
        math: { operator: "+" },
        format: { template: "{0}" },
        logic: { operator: "AND" },
        not: {},
        branch: {},
        sequence: {},
      };
      const newNode: Node = {
        id,
        type,
        position: pos,
        data: defaults[type] ?? {},
      };
      setNodes((nds) => [...nds, newNode]);
    },
    [setNodes],
  );

  const buttonClass =
    "px-2 py-1 text-[10px] font-medium rounded border transition-colors cursor-pointer";

  const toolbarButtons = (
    <div className="flex flex-wrap items-center gap-1 px-2 py-1.5 bg-zinc-900 rounded-t-md border border-zinc-700">
      <span className="text-[9px] uppercase tracking-wider text-zinc-500 font-semibold mr-1">Story</span>
      <button className={`${buttonClass} bg-blue-950 border-blue-700 text-blue-200 hover:bg-blue-900`} onClick={() => addNode("narrativeBeat")}>Beat</button>
      <button className={`${buttonClass} bg-amber-950 border-amber-700 text-amber-200 hover:bg-amber-900`} onClick={() => addNode("choice")}>Choice</button>
      <button className={`${buttonClass} bg-emerald-950 border-emerald-700 text-emerald-200 hover:bg-emerald-900`} onClick={() => addNode("outcome")}>Outcome</button>
      <button className={`${buttonClass} bg-yellow-950 border-yellow-700 text-yellow-200 hover:bg-yellow-900`} onClick={() => addNode("wait")}>Wait</button>

      <span className="text-[9px] uppercase tracking-wider text-zinc-500 font-semibold ml-2 mr-1">Data</span>
      <button className={`${buttonClass} bg-purple-950 border-purple-700 text-purple-200 hover:bg-purple-900`} onClick={() => addNode("number")}>Number</button>
      <button className={`${buttonClass} bg-teal-950 border-teal-700 text-teal-200 hover:bg-teal-900`} onClick={() => addNode("string")}>String</button>
      <button className={`${buttonClass} bg-emerald-950/30 border-emerald-700/50 text-emerald-200 hover:bg-emerald-900/40`} onClick={() => addNode("boolean")}>Bool</button>
      <button className={`${buttonClass} bg-cyan-950 border-cyan-700 text-cyan-200 hover:bg-cyan-900`} onClick={() => addNode("date")}>Date</button>
      <button className={`${buttonClass} bg-pink-950 border-pink-700 text-pink-200 hover:bg-pink-900`} onClick={() => addNode("random")}>Random</button>
      <button className={`${buttonClass} bg-indigo-950 border-indigo-700 text-indigo-200 hover:bg-indigo-900`} onClick={() => addNode("gameData")}>GameData</button>
      <button className={`${buttonClass} bg-orange-950 border-orange-700 text-orange-200 hover:bg-orange-900`} onClick={() => addNode("flag")}>Flag</button>

      <span className="text-[9px] uppercase tracking-wider text-zinc-500 font-semibold ml-2 mr-1">Logic</span>
      <button className={`${buttonClass} bg-slate-950 border-slate-700 text-slate-200 hover:bg-slate-900`} onClick={() => addNode("compare")}>Compare</button>
      <button className={`${buttonClass} bg-violet-950 border-violet-700 text-violet-200 hover:bg-violet-900`} onClick={() => addNode("math")}>Math</button>
      <button className={`${buttonClass} bg-sky-950 border-sky-700 text-sky-200 hover:bg-sky-900`} onClick={() => addNode("format")}>Format</button>
      <button className={`${buttonClass} bg-lime-950 border-lime-700 text-lime-200 hover:bg-lime-900`} onClick={() => addNode("logic")}>Logic</button>
      <button className={`${buttonClass} bg-red-950 border-red-700 text-red-200 hover:bg-red-900`} onClick={() => addNode("not")}>NOT</button>
      <button className={`${buttonClass} bg-rose-950 border-rose-700 text-rose-200 hover:bg-rose-900`} onClick={() => addNode("branch")}>Branch</button>
      <button className={`${buttonClass} bg-fuchsia-950 border-fuchsia-700 text-fuchsia-200 hover:bg-fuchsia-900`} onClick={() => addNode("sequence")}>Seq</button>

      <span className="ml-auto text-[10px] text-zinc-500">
        {nodes.length} node{nodes.length !== 1 ? "s" : ""}, {edges.length} edge{edges.length !== 1 ? "s" : ""}
      </span>
    </div>
  );

  return (
    <div className="flex flex-col gap-2 w-full" style={{ height: 600 }}>
      {toolbarButtons}
      <div className="flex-1 border border-zinc-700 rounded-b-md overflow-hidden">
        <ReactFlow
          nodes={nodes}
          edges={edges}
          onNodesChange={onNodesChange}
          onEdgesChange={onEdgesChange}
          onConnect={onConnect}
          isValidConnection={isValidConnection}
          onNodeClick={onNodeClick}
          onNodeContextMenu={onNodeContextMenu}
          onEdgeContextMenu={onEdgeContextMenu}
          onPaneContextMenu={onPaneContextMenu}
          nodeTypes={nodeTypes}
          fitView
          colorMode="dark"
          defaultEdgeOptions={{ animated: true }}
        >
          <Background color="#333" gap={20} />
          <Controls className="!bg-zinc-900 !border-zinc-700" />
          <MiniMap
            className="!border-zinc-700"
            style={{ background: "#18181b" }}
            nodeColor={(n) => {
              const colors: Record<string, string> = {
                narrativeBeat: "#3b82f6",
                choice: "#d97706",
                outcome: "#10b981",
                wait: "#facc15",
                number: "#a855f7",
                string: "#14b8a6",
                random: "#ec4899",
                boolean: "#34d399",
                date: "#06b6d4",
                gameData: "#6366f1",
                flag: "#f97316",
                compare: "#64748b",
                math: "#8b5cf6",
                format: "#0ea5e9",
                logic: "#84cc16",
                not: "#ef4444",
                branch: "#f43f5e",
                sequence: "#d946ef",
              };
              return colors[n.type as string] ?? "#6366f1";
            }}
          />
        </ReactFlow>

        {contextMenu && (
          <ContextMenuPrimitive.Root
            open
            onOpenChange={(open) => {
              if (!open) setContextMenu(null);
            }}
          >
            <ContextMenuPrimitive.Portal>
              <ContextMenuPrimitive.Content
                style={{ left: contextMenu.x, top: contextMenu.y, position: "fixed" }}
                onInteractOutside={() => setContextMenu(null)}
                className="z-50 min-w-[8rem] overflow-hidden rounded-md border border-zinc-700 bg-zinc-900 p-1 shadow-lg"
              >
                <ContextMenuPrimitive.Item
                  onClick={deleteSelected}
                  className="relative flex cursor-pointer items-center gap-2 rounded-sm px-2 py-1.5 text-sm text-red-400 outline-none select-none hover:bg-red-950/40"
                >
                  <Trash2Icon className="size-4" />
                  Delete{contextMenu.nodeId ? " node" : contextMenu.edgeId ? " connection" : ""}
                </ContextMenuPrimitive.Item>
              </ContextMenuPrimitive.Content>
            </ContextMenuPrimitive.Portal>
          </ContextMenuPrimitive.Root>
        )}

        <NodeEditDialog
          node={selectedNode}
          open={editingNodeId !== null}
          onOpenChange={(open) => {
            if (!open) setEditingNodeId(null);
          }}
          onSave={handleNodeSave}
        />
      </div>
    </div>
  );
}

export default NarrativeFlowEditor;