import type { ExecutionViewV8 } from "@/lib/types";

export interface ExpansionState {
  rounds: Set<number>;
  branches: Set<string>;
  retrievalAgents: Set<string>;
}

export function expandedStateFor(view: ExecutionViewV8): ExpansionState {
  return {
    rounds: new Set(view.rounds.map((round) => round.roundIndex)),
    branches: new Set(
      view.rounds.flatMap((round) =>
        round.branches.map((branch) => `${round.roundIndex}:${branch.branchId}`),
      ),
    ),
    retrievalAgents: new Set(
      view.rounds.flatMap((round) =>
        round.branches.flatMap((branch) =>
          branch.retrievalStage.agentRuns.map((agent) => agent.id),
        ),
      ),
    ),
  };
}

export function isFullyExpanded(
  view: ExecutionViewV8,
  expandedRounds: Set<number>,
  expandedBranches: Set<string>,
  expandedRetrievalAgents: Set<string>,
): boolean {
  const expected = expandedStateFor(view);
  return (
    [...expected.rounds].every((id) => expandedRounds.has(id))
    && [...expected.branches].every((id) => expandedBranches.has(id))
    && [...expected.retrievalAgents].every((id) => expandedRetrievalAgents.has(id))
  );
}