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 { NodeEditDialog } from "./NodeEditDialog";

const nodeTypes = {
  narrativeBeat: NarrativeBeatNode,
  choice: ChoiceNode,
  outcome: OutcomeNode,
};

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 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") => {
      const id = crypto.randomUUID();
      const pos = { x: Math.random() * 400, y: Math.random() * 300 };
      const newNode: Node = {
        id,
        type,
        position: pos,
        data:
          type === "choice"
            ? { question: "New choice?", options: [] }
            : type === "outcome"
              ? { title: "New outcome", text: "", effects: [] }
              : { title: "New beat", text: "" },
      };
      setNodes((nds) => [...nds, newNode]);
    },
    [setNodes],
  );

  const buttonClass =
    "px-3 py-1.5 text-xs font-medium rounded-md border transition-colors cursor-pointer";

  return (
    <div className="flex flex-col gap-2 w-full" style={{ height: 600 }}>
      <div className="flex items-center gap-2 px-2 py-1.5 bg-zinc-900 rounded-t-md border border-zinc-700">
        <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>
        <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>
      <div className="flex-1 border border-zinc-700 rounded-b-md overflow-hidden">
        <ReactFlow
          nodes={nodes}
          edges={edges}
          onNodesChange={onNodesChange}
          onEdgesChange={onEdgesChange}
          onConnect={onConnect}
          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) =>
              n.type === "choice"
                ? "#d97706"
                : n.type === "outcome"
                  ? "#10b981"
                  : "#3b82f6"
            }
          />
        </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;