EmboFlow/apps/web/src/runtime/workflow-editor-state.ts

export type WorkflowLogicNode = {
  id: string;
  type: string;
};

export type WorkflowLogicEdge = {
  from: string;
  to: string;
};

export type WorkflowPoint = {
  x: number;
  y: number;
};

export type WorkflowViewport = WorkflowPoint & {
  zoom: number;
};

export type WorkflowVisualGraph = {
  viewport: WorkflowViewport;
  nodePositions: Record<string, WorkflowPoint>;
};

export type WorkflowNodeDefinitionSummary = {
  id: string;
  name: string;
  category?: string;
  defaultExecutorType?: "python" | "docker" | "http";
  defaultExecutorConfig?: Record<string, unknown>;
  allowsMultipleIncoming?: boolean;
  supportsCodeHook?: boolean;
};

export type WorkflowCodeHookSpec = {
  language: "python";
  entrypoint?: string;
  source: string;
};

export type WorkflowNodeRuntimeConfig = {
  definitionId?: string;
  executorType?: "python" | "docker" | "http";
  executorConfig?: Record<string, unknown>;
  codeHookSpec?: WorkflowCodeHookSpec;
  artifactType?: "json" | "directory" | "video";
  artifactTitle?: string;
};

export type WorkflowDraft = {
  visualGraph: WorkflowVisualGraph;
  logicGraph: {
    nodes: WorkflowLogicNode[];
    edges: WorkflowLogicEdge[];
  };
  runtimeGraph: Record<string, unknown> & {
    selectedPreset?: string;
    nodeBindings?: Record<string, string>;
    nodeConfigs?: Record<string, WorkflowNodeRuntimeConfig>;
  };
  pluginRefs: string[];
};

type WorkflowVersionLike = Partial<WorkflowDraft>;

export type WorkflowConnectionValidationReason =
  | "missing_source"
  | "missing_target"
  | "self"
  | "duplicate"
  | "source_disallows_outgoing"
  | "target_disallows_incoming"
  | "target_already_has_incoming"
  | "cycle";

export type WorkflowConnectionValidationResult =
  | { ok: true }
  | { ok: false; reason: WorkflowConnectionValidationReason };

const DEFAULT_VIEWPORT: WorkflowViewport = { x: 0, y: 0, zoom: 1 };
const DEFAULT_NODE_LAYOUT: Record<string, WorkflowPoint> = {
  "source-asset": { x: 120, y: 120 },
  "rename-folder": { x: 430, y: 280 },
  "validate-structure": { x: 760, y: 450 },
};
const MULTI_INPUT_NODE_DEFINITION_IDS = new Set(["union-assets", "intersect-assets", "difference-assets"]);

function createDefaultNodePosition(index: number): WorkflowPoint {
  const column = index % 3;
  const row = Math.floor(index / 3);
  return {
    x: 140 + column * 280,
    y: 120 + row * 180,
  };
}

function createVisualGraph(input?: Partial<WorkflowVisualGraph>, nodes: WorkflowLogicNode[] = []): WorkflowVisualGraph {
  const nodePositions = { ...(input?.nodePositions ?? {}) };
  nodes.forEach((node, index) => {
    nodePositions[node.id] ??= DEFAULT_NODE_LAYOUT[node.id] ?? createDefaultNodePosition(index);
  });
  return {
    viewport: input?.viewport ? { ...DEFAULT_VIEWPORT, ...input.viewport } : { ...DEFAULT_VIEWPORT },
    nodePositions,
  };
}

function cloneDraft(draft: WorkflowDraft): WorkflowDraft {
  return {
    visualGraph: {
      viewport: { ...draft.visualGraph.viewport },
      nodePositions: Object.fromEntries(
        Object.entries(draft.visualGraph.nodePositions).map(([nodeId, position]) => [
          nodeId,
          { ...position },
        ]),
      ),
    },
    logicGraph: {
      nodes: draft.logicGraph.nodes.map((node) => ({ ...node })),
      edges: draft.logicGraph.edges.map((edge) => ({ ...edge })),
    },
    runtimeGraph: {
      ...draft.runtimeGraph,
      nodeBindings: { ...(draft.runtimeGraph.nodeBindings ?? {}) },
      nodeConfigs: Object.fromEntries(
        Object.entries(draft.runtimeGraph.nodeConfigs ?? {}).map(([nodeId, config]) => [
          nodeId,
          {
            ...config,
            executorConfig: config.executorConfig ? { ...config.executorConfig } : undefined,
            codeHookSpec: config.codeHookSpec ? { ...config.codeHookSpec } : undefined,
          },
        ]),
      ),
    },
    pluginRefs: [...draft.pluginRefs],
  };
}

function inferNodeType(definition: WorkflowNodeDefinitionSummary): string {
  const category = definition.category?.toLowerCase();
  if (category === "source" || category === "transform" || category === "inspect" || category === "export") {
    return category;
  }

  if (definition.id.startsWith("source")) {
    return "source";
  }
  if (
    definition.id.startsWith("extract") ||
    definition.id.startsWith("rename") ||
    definition.id.startsWith("transform")
  ) {
    return "transform";
  }
  if (definition.id.startsWith("validate") || definition.id.startsWith("inspect")) {
    return "inspect";
  }
  if (definition.id.startsWith("export")) {
    return "export";
  }
  return "utility";
}

function inferDefinitionId(nodeId: string): string {
  return nodeId.replace(/-\d+$/, "");
}

export function createDefaultWorkflowDraft(): WorkflowDraft {
  const logicGraph = {
    nodes: [
      { id: "source-asset", type: "source" },
      { id: "rename-folder", type: "transform" },
      { id: "validate-structure", type: "inspect" },
    ],
    edges: [
      { from: "source-asset", to: "rename-folder" },
      { from: "rename-folder", to: "validate-structure" },
    ],
  };
  return {
    visualGraph: createVisualGraph(undefined, logicGraph.nodes),
    logicGraph,
    runtimeGraph: {
      selectedPreset: "delivery-normalization",
      nodeBindings: {
        "source-asset": "source-asset",
        "rename-folder": "rename-folder",
        "validate-structure": "validate-structure",
      },
      nodeConfigs: {},
    },
    pluginRefs: ["builtin:delivery-nodes"],
  };
}

export function workflowDraftFromVersion(version?: WorkflowVersionLike | null): WorkflowDraft {
  if (!version?.logicGraph?.nodes?.length) {
    return createDefaultWorkflowDraft();
  }

  const nodeBindings = {
    ...(version.runtimeGraph?.nodeBindings ?? {}),
  } as Record<string, string>;
  const nodeConfigs = Object.fromEntries(
    Object.entries(version.runtimeGraph?.nodeConfigs ?? {}).map(([nodeId, config]) => [
      nodeId,
      {
        ...(config as WorkflowNodeRuntimeConfig),
        executorConfig: (config as WorkflowNodeRuntimeConfig)?.executorConfig
          ? { ...(config as WorkflowNodeRuntimeConfig).executorConfig }
          : undefined,
        codeHookSpec: (config as WorkflowNodeRuntimeConfig)?.codeHookSpec
          ? { ...(config as WorkflowNodeRuntimeConfig).codeHookSpec }
          : undefined,
      },
    ]),
  ) as Record<string, WorkflowNodeRuntimeConfig>;

  for (const node of version.logicGraph.nodes) {
    nodeBindings[node.id] ??= inferDefinitionId(node.id);
  }

  return {
    visualGraph: createVisualGraph(version.visualGraph, version.logicGraph.nodes),
    logicGraph: {
      nodes: version.logicGraph.nodes.map((node) => ({ ...node })),
      edges: (version.logicGraph.edges ?? []).map((edge) => ({ ...edge })),
    },
    runtimeGraph: {
      ...(version.runtimeGraph ?? {}),
      nodeBindings,
      nodeConfigs,
    },
    pluginRefs: [...(version.pluginRefs ?? [])],
  };
}

export function addNodeToDraft(
  draft: WorkflowDraft,
  definition: WorkflowNodeDefinitionSummary,
): { draft: WorkflowDraft; nodeId: string } {
  return addNodeToDraftInternal(draft, definition, {
    connectFromPrevious: true,
  });
}

export function addNodeToDraftAtPosition(
  draft: WorkflowDraft,
  definition: WorkflowNodeDefinitionSummary,
  position: WorkflowPoint,
): { draft: WorkflowDraft; nodeId: string } {
  return addNodeToDraftInternal(draft, definition, {
    connectFromPrevious: false,
    position,
  });
}

function addNodeToDraftInternal(
  draft: WorkflowDraft,
  definition: WorkflowNodeDefinitionSummary,
  options: {
    connectFromPrevious: boolean;
    position?: WorkflowPoint;
  },
): { draft: WorkflowDraft; nodeId: string } {
  const next = cloneDraft(draft);
  let suffix = 1;
  let nodeId = `${definition.id}-${suffix}`;
  const existingIds = new Set(next.logicGraph.nodes.map((node) => node.id));
  while (existingIds.has(nodeId)) {
    suffix += 1;
    nodeId = `${definition.id}-${suffix}`;
  }

  const node: WorkflowLogicNode = {
    id: nodeId,
    type: inferNodeType(definition),
  };
  const previousNode = next.logicGraph.nodes.at(-1);
  next.logicGraph.nodes.push(node);
  if (options.connectFromPrevious && previousNode) {
    next.logicGraph.edges.push({ from: previousNode.id, to: nodeId });
  }
  next.runtimeGraph.nodeBindings ??= {};
  next.runtimeGraph.nodeBindings[nodeId] = definition.id;
  next.runtimeGraph.nodeConfigs ??= {};
  next.visualGraph.nodePositions[nodeId] =
    options.position ?? createDefaultNodePosition(next.logicGraph.nodes.length - 1);

  return { draft: next, nodeId };
}

export function removeNodeFromDraft(draft: WorkflowDraft, nodeId: string): WorkflowDraft {
  const next = cloneDraft(draft);
  next.logicGraph.nodes = next.logicGraph.nodes.filter((node) => node.id !== nodeId);
  next.logicGraph.edges = next.logicGraph.edges.filter(
    (edge) => edge.from !== nodeId && edge.to !== nodeId,
  );
  if (next.runtimeGraph.nodeBindings) {
    delete next.runtimeGraph.nodeBindings[nodeId];
  }
  if (next.runtimeGraph.nodeConfigs) {
    delete next.runtimeGraph.nodeConfigs[nodeId];
  }
  delete next.visualGraph.nodePositions[nodeId];
  return next;
}

export function resolveDefinitionIdForNode(draft: WorkflowDraft, nodeId: string): string {
  return draft.runtimeGraph.nodeBindings?.[nodeId] ?? inferDefinitionId(nodeId);
}

export function getNodeRuntimeConfig(draft: WorkflowDraft, nodeId: string) {
  return draft.runtimeGraph.nodeConfigs?.[nodeId];
}

export function setNodeRuntimeConfig(
  draft: WorkflowDraft,
  nodeId: string,
  config: WorkflowNodeRuntimeConfig,
): WorkflowDraft {
  const next = cloneDraft(draft);
  next.runtimeGraph.nodeConfigs ??= {};
  next.runtimeGraph.nodeConfigs[nodeId] = {
    ...config,
    executorConfig: config.executorConfig ? { ...config.executorConfig } : undefined,
    codeHookSpec: config.codeHookSpec ? { ...config.codeHookSpec } : undefined,
  };
  return next;
}

export function setNodePosition(draft: WorkflowDraft, nodeId: string, position: WorkflowPoint): WorkflowDraft {
  const next = cloneDraft(draft);
  next.visualGraph.nodePositions[nodeId] = { ...position };
  return next;
}

function findNode(draft: WorkflowDraft, nodeId: string) {
  return draft.logicGraph.nodes.find((node) => node.id === nodeId) ?? null;
}

function nodeDisallowsOutgoing(node: WorkflowLogicNode) {
  return node.type === "export";
}

function nodeDisallowsIncoming(node: WorkflowLogicNode) {
  return node.type === "source";
}

function allowsMultipleIncoming(draft: WorkflowDraft, nodeId: string) {
  const definitionId = resolveDefinitionIdForNode(draft, nodeId);
  return MULTI_INPUT_NODE_DEFINITION_IDS.has(definitionId);
}

function wouldCreateCycle(draft: WorkflowDraft, sourceNodeId: string, targetNodeId: string) {
  const adjacency = new Map<string, string[]>();
  for (const edge of draft.logicGraph.edges) {
    const current = adjacency.get(edge.from) ?? [];
    current.push(edge.to);
    adjacency.set(edge.from, current);
  }

  const stack = [targetNodeId];
  const visited = new Set<string>();
  while (stack.length > 0) {
    const currentNodeId = stack.pop();
    if (!currentNodeId || visited.has(currentNodeId)) {
      continue;
    }
    if (currentNodeId === sourceNodeId) {
      return true;
    }
    visited.add(currentNodeId);
    for (const nextNodeId of adjacency.get(currentNodeId) ?? []) {
      stack.push(nextNodeId);
    }
  }

  return false;
}

export function canConnectNodesInDraft(
  draft: WorkflowDraft,
  sourceNodeId: string | null | undefined,
  targetNodeId: string | null | undefined,
): WorkflowConnectionValidationResult {
  if (!sourceNodeId) {
    return { ok: false, reason: "missing_source" };
  }
  if (!targetNodeId) {
    return { ok: false, reason: "missing_target" };
  }
  if (sourceNodeId === targetNodeId) {
    return { ok: false, reason: "self" };
  }

  const sourceNode = findNode(draft, sourceNodeId);
  if (!sourceNode) {
    return { ok: false, reason: "missing_source" };
  }

  const targetNode = findNode(draft, targetNodeId);
  if (!targetNode) {
    return { ok: false, reason: "missing_target" };
  }

  if (draft.logicGraph.edges.some((edge) => edge.from === sourceNodeId && edge.to === targetNodeId)) {
    return { ok: false, reason: "duplicate" };
  }
  if (nodeDisallowsOutgoing(sourceNode)) {
    return { ok: false, reason: "source_disallows_outgoing" };
  }
  if (nodeDisallowsIncoming(targetNode)) {
    return { ok: false, reason: "target_disallows_incoming" };
  }
  if (!allowsMultipleIncoming(draft, targetNodeId) && draft.logicGraph.edges.some((edge) => edge.to === targetNodeId)) {
    return { ok: false, reason: "target_already_has_incoming" };
  }
  if (wouldCreateCycle(draft, sourceNodeId, targetNodeId)) {
    return { ok: false, reason: "cycle" };
  }

  return { ok: true };
}

export function connectNodesInDraft(
  draft: WorkflowDraft,
  sourceNodeId: string | null | undefined,
  targetNodeId: string | null | undefined,
): WorkflowDraft {
  const validation = canConnectNodesInDraft(draft, sourceNodeId, targetNodeId);
  if (!validation.ok || !sourceNodeId || !targetNodeId) {
    return draft;
  }
  const next = cloneDraft(draft);
  next.logicGraph.edges.push({ from: sourceNodeId, to: targetNodeId });
  return next;
}

export function setViewportInDraft(draft: WorkflowDraft, viewport: WorkflowViewport): WorkflowDraft {
  const next = cloneDraft(draft);
  next.visualGraph.viewport = { ...viewport };
  return next;
}

export function serializeWorkflowDraft(draft: WorkflowDraft): WorkflowDraft {
  return cloneDraft(draft);
}