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 { 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,
};

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);
      if (!sourceSocket || !targetSocket) return false;

      if (sourceSocket.type === "any" || targetSocket.type === "any") return true;
      return sourceSocket.type === targetSocket.type;
    },
    [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",
    ) => {
      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: [] }
              : type === "wait"
                ? { ticks: 10 }
                : type === "number"
                  ? { value: 0 }
                  : type === "string"
                    ? { value: "" }
                    : type === "random"
                      ? { min: 0, max: 100 }
                      : type === "boolean"
                        ? { value: false }
                        : type === "date"
                          ? { value: "" }
                          : { 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";

  const toolbarButtons = (
    <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>
      <button
        className={`${buttonClass} bg-yellow-950 border-yellow-700 text-yellow-200 hover:bg-yellow-900`}
        onClick={() => addNode("wait")}
      >
        + Wait
      </button>
      <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-pink-950 border-pink-700 text-pink-200 hover:bg-pink-900`}
        onClick={() => addNode("random")}
      >
        + Random
      </button>
      <button
        className={`${buttonClass} bg-emerald-950/30 border-emerald-700/50 text-emerald-200 hover:bg-emerald-900/40`}
        onClick={() => addNode("boolean")}
      >
        + Boolean
      </button>
      <button
        className={`${buttonClass} bg-cyan-950 border-cyan-700 text-cyan-200 hover:bg-cyan-900`}
        onClick={() => addNode("date")}
      >
        + Date
      </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) => {
              switch (n.type) {
                case "choice":
                  return "#d97706"; // amber
                case "outcome":
                  return "#10b981"; // emerald
                case "narrativeBeat":
                  return "#3b82f6"; // blue
                case "wait":
                  return "#facc15"; // yellow
                case "number":
                  return "#a855f7"; // purple
                case "string":
                  return "#14b8a6"; // teal
                case "random":
                  return "#ec4899"; // pink
                case "boolean":
                  return "#34d399"; // light green
                case "date":
                  return "#06b6d4"; // cyan
                default:
                  return "#6366f1"; // indigo
              }
            }}
          />
        </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;