import type {
  ChildTaskId,
  ChildTaskRecord,
  ChildTaskStatus,
  TurnId,
  VerdictRecord,
  WorkflowName,
  WorkflowNode,
  WorkflowPhaseName,
} from "@best-agent/protocol";

export type WorkflowGraphEdge = {
  from: WorkflowPhaseName;
  to: WorkflowPhaseName;
};

export function resolveWorkflowPhaseLayers(
  workflowName: WorkflowName,
): Array<Array<WorkflowPhaseName>> {
  switch (workflowName) {
    case "coordinate":
      return [];
    case "explore-build":
      return [["explore"], ["build"]];
    case "build-verify":
      return [["build"], ["verify"]];
    case "plan-build-verify":
      return [["plan"], ["build"], ["verify"]];
    case "plan-explore-build-verify":
      return [["plan", "explore"], ["build"], ["verify"]];
    default:
      return [["plan"], ["build"], ["verify"]];
  }
}

export function inferWorkflowNameFromPhaseResults(
  phaseResults: Partial<Record<WorkflowPhaseName, string>>,
): WorkflowName {
  if (phaseResults.explore && phaseResults.plan && phaseResults.build && phaseResults.verify) {
    return "plan-explore-build-verify";
  }
  if (phaseResults.explore && phaseResults.build && !phaseResults.plan && !phaseResults.verify) {
    return "explore-build";
  }
  if (phaseResults.build && phaseResults.verify && !phaseResults.explore && !phaseResults.plan) {
    return "build-verify";
  }
  return "plan-build-verify";
}

export function buildWorkflowPhaseDependencies(
  workflowLayers: Array<Array<WorkflowPhaseName>>,
): Map<WorkflowPhaseName, WorkflowPhaseName[]> {
  const dependencies = new Map<WorkflowPhaseName, WorkflowPhaseName[]>();
  let priorLayer: WorkflowPhaseName[] = [];
  for (const layer of workflowLayers) {
    for (const phaseName of layer) {
      dependencies.set(phaseName, [...priorLayer]);
    }
    priorLayer = layer;
  }
  return dependencies;
}

export function buildWorkflowPhaseEdges(
  workflowLayers: Array<Array<WorkflowPhaseName>>,
): WorkflowGraphEdge[] {
  return workflowLayers.flatMap((layer, index) => {
    const next = workflowLayers[index + 1];
    if (!next) {
      return [];
    }
    return layer.flatMap((from) => next.map((to) => ({ from, to })));
  });
}

export function resolveReadyWorkflowPhases(input: {
  completedPhases: ReadonlySet<WorkflowPhaseName>;
  pendingPhases: ReadonlySet<WorkflowPhaseName>;
  workflowLayers: Array<Array<WorkflowPhaseName>>;
}): WorkflowPhaseName[] {
  const workflowDependencies = buildWorkflowPhaseDependencies(input.workflowLayers);
  return input.workflowLayers.flatMap((layer) =>
    layer.filter((phaseName) => {
      if (!input.pendingPhases.has(phaseName)) {
        return false;
      }
      const dependencies = workflowDependencies.get(phaseName) ?? [];
      return dependencies.every((dependency) => input.completedPhases.has(dependency));
    }),
  );
}

export function resolveWorkflowNodeAllowedTools(input: {
  fallbackPhaseName: WorkflowPhaseName;
  profileName?: string;
  readProfileAllowedTools: (profileName: string) => string[] | undefined;
}): string[] {
  const profileTools = input.profileName
    ? input.readProfileAllowedTools(input.profileName)
    : undefined;
  if (profileTools) {
    return [...profileTools];
  }
  return [...(input.readProfileAllowedTools(input.fallbackPhaseName) ?? [])];
}

export function buildWorkflowChildTaskGraphNodes(input: {
  effectiveVerdicts?: ReadonlyMap<string, VerdictRecord>;
  phases: ReadonlyMap<
    WorkflowPhaseName,
    ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
  >;
  readProfileAllowedTools: (profileName: string) => string[] | undefined;
  statusForRecord?: (
    record:
      | (ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName })
      | undefined,
  ) => WorkflowNode["status"];
  workflowName: WorkflowName;
}): WorkflowNode[] {
  const steps = resolveWorkflowPhaseLayers(input.workflowName);
  return steps.flatMap((step) =>
    step.map((phaseName) => {
      const record = input.phases.get(phaseName);
      const profileName = record?.profileName ?? phaseName;
      return {
        allowedTools: resolveWorkflowNodeAllowedTools({
          fallbackPhaseName: phaseName,
          profileName,
          readProfileAllowedTools: input.readProfileAllowedTools,
        }),
        ...(record?.childThreadId ? { childThreadId: record.childThreadId } : {}),
        ...(record?.correlationId ? { correlationId: record.correlationId } : {}),
        ...(record?.childTaskId && input.effectiveVerdicts?.get(record.childTaskId)
          ? { effectiveVerdict: input.effectiveVerdicts.get(record.childTaskId) }
          : {}),
        ...(record?.childTaskId ? { nodeId: record.childTaskId } : {}),
        phaseName,
        profileName,
        status: input.statusForRecord
          ? input.statusForRecord(record)
          : (record?.status ?? "pending"),
      };
    }),
  );
}

export function buildWorkflowChildTaskGraph(input: {
  effectiveVerdicts?: ReadonlyMap<string, VerdictRecord>;
  phases: ReadonlyMap<
    WorkflowPhaseName,
    ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
  >;
  readProfileAllowedTools: (profileName: string) => string[] | undefined;
  statusForRecord?: (
    record:
      | (ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName })
      | undefined,
  ) => WorkflowNode["status"];
  workflowName: WorkflowName;
}): {
  edges: WorkflowGraphEdge[];
  nodes: WorkflowNode[];
} {
  return {
    edges: buildWorkflowPhaseEdges(resolveWorkflowPhaseLayers(input.workflowName)),
    nodes: buildWorkflowChildTaskGraphNodes(input),
  };
}

export function buildWorkflowPhaseResultGraph(input: {
  phaseResults: Partial<Record<WorkflowPhaseName, string>>;
  readProfileAllowedTools: (profileName: string) => string[] | undefined;
  workflowName: WorkflowName;
}): {
  edges: WorkflowGraphEdge[];
  nodes: WorkflowNode[];
} {
  const workflowLayers = resolveWorkflowPhaseLayers(input.workflowName);
  return {
    edges: buildWorkflowPhaseEdges(workflowLayers),
    nodes: workflowLayers.flatMap((layer) =>
      layer.map((phaseName) => ({
        allowedTools: resolveWorkflowNodeAllowedTools({
          fallbackPhaseName: phaseName,
          readProfileAllowedTools: input.readProfileAllowedTools,
        }),
        phaseName,
        profileName: phaseName,
        status: input.phaseResults[phaseName] ? "completed" : "pending",
      })),
    ),
  };
}

export function groupWorkflowPhaseRecordsByTurn(
  records: ReadonlyArray<ChildTaskRecord>,
): Map<
  TurnId,
  Map<
    WorkflowPhaseName,
    ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
  >
> {
  const latestByTurn = new Map<
    TurnId,
    Map<
      WorkflowPhaseName,
      ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
    >
  >();
  for (const record of records) {
    if (!record.phaseName || !record.workflowName) {
      continue;
    }
    let phases = latestByTurn.get(record.turnId);
    if (!phases) {
      phases = new Map();
      latestByTurn.set(record.turnId, phases);
    }
    const existing = phases.get(record.phaseName);
    if (!existing || shouldReplaceWorkflowPhaseRecord(existing, record)) {
      phases.set(
        record.phaseName,
        record as ChildTaskRecord & {
          phaseName: WorkflowPhaseName;
          workflowName: WorkflowName;
        },
      );
    }
  }
  return latestByTurn;
}

function shouldReplaceWorkflowPhaseRecord(
  existing: ChildTaskRecord,
  candidate: ChildTaskRecord,
): boolean {
  if (existing.childTaskId === candidate.childTaskId) {
    return childTaskStatusRank(existing.status) <= childTaskStatusRank(candidate.status);
  }
  return existing.createdAt <= candidate.createdAt;
}

function childTaskStatusRank(status: ChildTaskStatus): number {
  switch (status) {
    case "running":
      return 0;
    case "completed":
    case "failed":
      return 1;
    case "joined":
    case "rejected":
      return 2;
  }
}

export function listWorkflowVerifyChildTaskIds(
  phases: ReadonlyMap<
    WorkflowPhaseName,
    ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
  >,
): ChildTaskId[] {
  return Array.from(phases.values())
    .filter((record) => record.phaseName === "verify")
    .map((record) => record.childTaskId);
}

export async function collectChildTaskEffectiveVerdicts(input: {
  childTaskIds: ReadonlyArray<ChildTaskId>;
  effectiveVerdicts: ReadonlyMap<string, VerdictRecord>;
}): Promise<Map<string, VerdictRecord>> {
  const selectedVerdicts = new Map<string, VerdictRecord>();
  for (const nodeId of input.childTaskIds) {
    const verdict = input.effectiveVerdicts.get(nodeId);
    if (verdict) {
      selectedVerdicts.set(nodeId, verdict);
    }
  }
  return selectedVerdicts;
}

export async function collectWorkflowVerifyVerdicts(input: {
  effectiveVerdicts: ReadonlyMap<string, VerdictRecord>;
  phases: ReadonlyMap<
    WorkflowPhaseName,
    ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
  >;
}): Promise<Map<string, VerdictRecord>> {
  return collectChildTaskEffectiveVerdicts({
    childTaskIds: listWorkflowVerifyChildTaskIds(input.phases),
    effectiveVerdicts: input.effectiveVerdicts,
  });
}

export async function collectWorkflowTurnStates(input: {
  effectiveVerdicts: ReadonlyMap<string, VerdictRecord>;
  records: ReadonlyArray<ChildTaskRecord>;
}): Promise<
  Array<{
    effectiveVerdicts: Map<string, VerdictRecord>;
    phases: Map<
      WorkflowPhaseName,
      ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
    >;
    turnId: TurnId;
    workflowName: WorkflowName;
  }>
> {
  const latestByTurn = groupWorkflowPhaseRecordsByTurn(input.records);
  const workflowTurns: Array<{
    effectiveVerdicts: Map<string, VerdictRecord>;
    phases: Map<
      WorkflowPhaseName,
      ChildTaskRecord & { phaseName: WorkflowPhaseName; workflowName: WorkflowName }
    >;
    turnId: TurnId;
    workflowName: WorkflowName;
  }> = [];

  for (const [turnId, phases] of latestByTurn.entries()) {
    const workflowName = phases.values().next().value?.workflowName;
    if (!workflowName) {
      continue;
    }
    workflowTurns.push({
      effectiveVerdicts: await collectWorkflowVerifyVerdicts({
        effectiveVerdicts: input.effectiveVerdicts,
        phases,
      }),
      phases,
      turnId,
      workflowName,
    });
  }

  return workflowTurns;
}
