import type { BackendTask, ProcessArtifact, ProcessEvent, ProcessRun } from '@/types';

export type TaskProcessSelection = {
  runs: ProcessRun[];
  events: ProcessEvent[];
  artifacts: ProcessArtifact[];
};

export type SelectTaskProcessInput = {
  task: BackendTask;
  liveRuns?: ProcessRun[];
  liveEvents?: ProcessEvent[];
  liveArtifacts?: ProcessArtifact[];
};

function mergeById<T>(
  persisted: T[],
  live: T[],
  idFor: (item: T) => string,
): T[] {
  const merged = new Map<string, T>();
  for (const item of persisted) merged.set(idFor(item), item);
  for (const item of live) merged.set(idFor(item), item);
  return Array.from(merged.values());
}

function runBelongsToTask(run: ProcessRun, taskId: string, runIds: Set<string>): boolean {
  return runIds.has(run.run_id) || run.metadata?.task_id === taskId;
}

function expandTaskRunIds(runs: ProcessRun[], taskId: string, seedRunIds: Set<string>): Set<string> {
  const selected = new Set(seedRunIds);
  for (const run of runs) {
    if (run.metadata?.task_id === taskId) selected.add(run.run_id);
  }

  let changed = true;
  while (changed) {
    changed = false;
    for (const run of runs) {
      if (!run.parent_run_id || !selected.has(run.parent_run_id) || selected.has(run.run_id)) continue;
      selected.add(run.run_id);
      changed = true;
    }
  }
  return selected;
}

export function selectTaskProcess({
  task,
  liveRuns = [],
  liveEvents = [],
  liveArtifacts = [],
}: SelectTaskProcessInput): TaskProcessSelection {
  const persistedRuns = task.process_runs ?? [];
  const persistedEvents = task.process_events ?? [];
  const persistedArtifacts = task.process_artifacts ?? [];
  const allRuns = mergeById(persistedRuns, liveRuns, (run) => run.run_id);
  const seedRunIds = new Set([
    ...task.run_ids.filter(Boolean),
    ...persistedRuns.map((run) => run.run_id),
  ]);
  const taskRunIds = expandTaskRunIds(allRuns, task.task_id, seedRunIds);
  const runs = allRuns.filter((run) => runBelongsToTask(run, task.task_id, taskRunIds));

  return {
    runs,
    events: mergeById(persistedEvents, liveEvents, (event) => event.event_id).filter(
      (event) => taskRunIds.has(event.run_id) || event.metadata?.task_id === task.task_id
    ),
    artifacts: mergeById(persistedArtifacts, liveArtifacts, (artifact) => artifact.artifact_id).filter(
      (artifact) => taskRunIds.has(artifact.run_id) || artifact.metadata?.task_id === task.task_id
    ),
  };
}