import type { AgentRunSnapshot } from "./protocol";

export function finalImageResultNodeIds(
  snapshot: AgentRunSnapshot,
  nodes: Array<{ id: string; type: string }>,
) {
  const images = new Set(nodes.filter((node) => node.type === "image").map((node) => node.id));
  return finalResultNodeIds(snapshot).filter((id) => images.has(id));
}

export function finalResultNodeIds(snapshot: AgentRunSnapshot) {
  const consumed = new Set(
    snapshot.calls
      .filter(
        (call) => Array.isArray(call.output?.createdNodeIds) && call.output.createdNodeIds.length,
      )
      .flatMap((call) => sourceIds(call.input, call.output)),
  );
  return snapshot.resultNodeIds.filter((id) => !consumed.has(id));
}
function sourceIds(input: Record<string, unknown>, output?: Record<string, unknown>) {
  return [
    ...(Array.isArray(input.referenceNodeIds) ? input.referenceNodeIds : []),
    ...(Array.isArray(input.nodeIds) ? input.nodeIds : []),
    ...(Array.isArray(output?.outputs) ? output.outputs.map((item) => item?.sourceNodeId) : []),
  ].filter((id): id is string => typeof id === "string");
}
export function originalSourceNodeIds(snapshot: AgentRunSnapshot, nodeIds: string[]) {
  const visited = new Set<string>();
  const originals = new Set<string>();
  const visit = (id: string) => {
    if (visited.has(id)) return;
    visited.add(id);
    const producer = snapshot.calls.find(
      (call) =>
        Array.isArray(call.output?.createdNodeIds) && call.output.createdNodeIds.includes(id),
    );
    const sources = producer ? sourceIds(producer.input, producer.output) : [];
    if (sources.length) sources.forEach(visit);
    else if (!snapshot.resultNodeIds.includes(id)) originals.add(id);
  };
  nodeIds.forEach(visit);
  return [...originals];
}