import type { ProviderConfig, ProviderTurnActionPolicy, Turn } from "@best-agent/protocol";
import type { AgentProfile } from "../../agent/profile-registry.js";
import type { PromptDeliverableFailure } from "../../prompt/runtime-prompts.js";
import type {
  CompletedProviderToolCall,
  ProviderCorrectionTrace,
  ToolOutputCheckpointPayload,
  ToolResumeContext,
} from "../checkpoint-codec.js";
import { assertReplayToolOutputDeliverables } from "./provider-replay-artifact-assertion.js";
import type { ProviderReplayRepairRuntime } from "./provider-replay-context.js";
import { runReplayRepairDecision } from "./provider-replay-decision-runner.js";
import {
  buildReplayToolOutputRepairContexts,
  prepareReplayToolOutputRepair,
} from "./provider-replay-tool-output-repair-context.js";
import type { ProviderTurnBudget } from "./provider-turn-budget.js";
import { appendRuntimeObservation, type RuntimeObservation } from "./runtime-observation.js";
import type { RequiredToolConstraint } from "./tool-constraint-types.js";

export type ReplayToolOutputRepairPayload = ToolResumeContext &
  Partial<
    Pick<
      ToolOutputCheckpointPayload,
      "toolEnvelope" | "toolInputStableDigest" | "toolOutput" | "toolResultItem"
    >
  >;

export async function repairReplayToolOutputWithObservation(input: {
  completedToolCalls: CompletedProviderToolCall[];
  fixedReplayRoute: boolean;
  payload: ReplayToolOutputRepairPayload;
  repairActionPolicyOverride?: ProviderTurnActionPolicy;
  repairCode: ProviderCorrectionTrace["code"];
  repairMessages: ReturnType<typeof appendRuntimeObservation>;
  repairObservation: RuntimeObservation;
  repairRouteConfig: ProviderConfig;
  responseRepairDeliverableFailures?: PromptDeliverableFailure[];
  replayProfile: AgentProfile;
  replayProviderTurnBudget: ProviderTurnBudget;
  requiresDeliverableAssertion?: boolean;
  suppressStructuredControls?: boolean;
  runtime: ProviderReplayRepairRuntime;
  turn: Turn;
  unresolvedRequiredToolConstraint?: RequiredToolConstraint;
}): Promise<string> {
  const {
    completedToolCalls,
    fixedReplayRoute,
    payload,
    repairCode,
    repairMessages,
    repairObservation,
    repairRouteConfig,
    replayProfile,
    replayProviderTurnBudget,
    runtime,
    turn,
    unresolvedRequiredToolConstraint,
  } = input;
  const repair = prepareReplayToolOutputRepair({
    completedToolCalls,
    payload,
    repairActionPolicyOverride: input.repairActionPolicyOverride,
    repairCode,
    repairMessages,
    repairObservation,
    requiresDeliverableAssertion: input.requiresDeliverableAssertion,
    suppressStructuredControls: input.suppressStructuredControls,
    unresolvedRequiredToolConstraint,
  });
  const repairContexts = buildReplayToolOutputRepairContexts({
    completedToolCalls,
    fixedReplayRoute,
    payload,
    repairRequiredToolConstraint: repair.repairRequiredToolConstraint,
    repairRouteConfig,
    replayProfile,
    replayProviderTurnBudget,
  });
  const repairedText = await runReplayRepairDecision({
    actionPolicy: repair.repairActionPolicy,
    allowPhaseResult: false,
    allowPlanUpdate: false,
    artifactObligations: payload.artifactObligations ?? [],
    availableTools: payload.availableTools,
    baseMessages: repairMessages,
    cacheControl: payload.cacheControl,
    controls: {
      allowPhaseResult: false,
      allowPlanUpdate: false,
      ...(payload.phaseResultPhaseName
        ? { phaseResultPhaseName: payload.phaseResultPhaseName }
        : {}),
    },
    correctionOperation: "tool-followup",
    phaseResultPhaseName: payload.phaseResultPhaseName,
    providerContext: repairContexts.providerContext,
    repairCode,
    repairObservation,
    responseRepairDeliverableFailures: input.responseRepairDeliverableFailures,
    requiredToolConstraint: repair.repairRequiredToolConstraint,
    resolveContext: repairContexts.resolveContext,
    runtime,
    objective: payload.objective,
    systemPrompt: repair.repairSystemPrompt,
    toolsForTrace: payload.availableTools,
    toolStepCount: repairContexts.toolStepCount,
    traceSystemPrompt: payload.systemPrompt,
    turn,
    visibleTools: repair.visibleTools,
  });
  await assertReplayToolOutputDeliverables({
    payload,
    repairMessages,
    requiresDeliverableAssertion: repair.requiresDeliverableAssertion,
    runtime,
    turn,
  });
  return repairedText;
}
