'use client';

import * as React from 'react';
import type { FlowNode, FlowEdge } from '../serialization';
import type { WfExecutionView, NodeStatus } from '../types';
import type { WorkflowNodeData } from '../WorkflowNode';
import type { WorkflowEdgeData } from '../WorkflowEdge';

/**
 * Merge an optional {@link WfExecutionView} status overlay onto flow nodes and
 * edges. Each node's `data.status` is set from the matching node execution
 * (defaulting to "idle"); each edge's `data.sourceStatus` is set from its
 * source node so {@link WorkflowEdge} can animate live runs.
 *
 * Returns fresh node/edge arrays (referentially stable when inputs are) so
 * xyflow re-renders only on actual overlay changes — suitable for polling.
 */
export function useNodeStatusOverlay(
  nodes: FlowNode[],
  edges: FlowEdge[],
  execution?: WfExecutionView
): { nodes: FlowNode[]; edges: FlowEdge[] } {
  return React.useMemo(() => {
    const statusFor = (nodeId: string): NodeStatus =>
      execution?.nodes?.[nodeId]?.status ?? 'idle';

    const overlaidNodes: FlowNode[] = nodes.map((n) => {
      const status = statusFor(n.id);
      const data = n.data as WorkflowNodeData;
      if (data.status === status) return n;
      return { ...n, data: { ...data, status } };
    });

    const overlaidEdges: FlowEdge[] = edges.map((e) => {
      const sourceStatus = statusFor(e.source);
      const data = (e.data ?? { connectionType: 'main' }) as FlowEdge['data'] &
        WorkflowEdgeData;
      if (data.sourceStatus === sourceStatus) return e;
      return {
        ...e,
        data: { ...data, connectionType: data.connectionType ?? 'main', sourceStatus },
      };
    });

    return { nodes: overlaidNodes, edges: overlaidEdges };
  }, [nodes, edges, execution]);
}