import type {
  ArtifactObligation,
  ProviderTurnActionPolicy,
  SessionUsageOperation,
  Turn,
} from "@best-agent/protocol";
import type {
  ConversationMessage,
  ProviderCacheControl,
  ProviderToolOption,
} from "../../provider/model-provider.js";
import type { PromptDeliverableFailure } from "../../prompt/runtime-prompts.js";
import type { ProviderCorrectionTrace } from "../checkpoint-codec.js";
import type { ProviderDecisionResolutionControls } from "./provider-decision-resolution.js";
import type { ProviderReplayRepairRuntime } from "./provider-replay-context.js";
import type { ProviderTurnInvokeContext } from "./provider-turn-budget.js";
import { appendRuntimeObservation, type RuntimeObservation } from "./runtime-observation.js";
import type { RequiredToolConstraint } from "./tool-constraint-types.js";

export async function runReplayRepairDecision(input: {
  actionPolicy?: ProviderTurnActionPolicy;
  allowPhaseResult: boolean;
  allowPlanUpdate: boolean;
  artifactObligations: ArtifactObligation[];
  assertArtifacts?: boolean;
  assertTerminal?: boolean;
  availableTools: ProviderToolOption[];
  baseMessages: ConversationMessage[];
  cacheControl?: ProviderCacheControl;
  controls: ProviderDecisionResolutionControls;
  correctionOperation: SessionUsageOperation;
  phaseResultPhaseName?: ProviderDecisionResolutionControls["phaseResultPhaseName"];
  providerContext: ProviderTurnInvokeContext;
  repairCode: ProviderCorrectionTrace["code"];
  repairObservation: RuntimeObservation;
  requiredToolConstraint?: RequiredToolConstraint;
  responseRepairDeliverableFailures?: PromptDeliverableFailure[];
  resolveContext: ProviderTurnInvokeContext;
  runtime: Pick<
    ProviderReplayRepairRuntime,
    | "assertArtifactObligationsSatisfied"
    | "assertValidTerminalAssistantTextOrFallback"
    | "buildCorrectionTrace"
    | "provider"
    | "resolveProviderDecision"
  >;
  systemPrompt?: string;
  toolsForTrace?: ProviderToolOption[];
  toolStepCount: number;
  traceSystemPrompt?: string;
  turn: Turn;
  visibleTools: ProviderToolOption[];
  objective?: string;
}): Promise<string> {
  const repairMessages = appendRuntimeObservation(input.baseMessages, input.repairObservation);
  const correctionTrace = input.runtime.buildCorrectionTrace({
    code: input.repairCode,
    messages: input.baseMessages,
    operation: input.correctionOperation,
    phaseResultPhaseName: input.phaseResultPhaseName,
    systemPrompt: input.traceSystemPrompt ?? input.systemPrompt,
    tools: input.toolsForTrace,
  });
  const decision = await input.runtime.provider.decide(
    input.turn.threadId,
    input.turn.id,
    {
      allowPhaseResult: input.allowPhaseResult,
      allowPlanUpdate: input.allowPlanUpdate,
      ...(input.actionPolicy ? { actionPolicy: input.actionPolicy } : {}),
      cacheControl: input.cacheControl,
      messages: repairMessages,
      ...(input.phaseResultPhaseName ? { phaseResultPhaseName: input.phaseResultPhaseName } : {}),
      systemPrompt: input.systemPrompt,
      tools: input.visibleTools,
    },
    "repair",
    {
      ...input.providerContext,
      correctionTrace,
    },
  );
  const repairedText = await input.runtime.resolveProviderDecision({
    turn: input.turn,
    messages: repairMessages,
    systemPrompt: input.systemPrompt,
    decision,
    availableTools: input.availableTools,
    visibleTools: input.visibleTools,
    cacheControl: input.cacheControl,
    controls: input.controls,
    invokeContext: {
      ...input.resolveContext,
      correctionTrace,
    },
    operation: "repair",
    objective: input.objective,
    repairAttempt: 0,
    requiredToolConstraint: input.requiredToolConstraint,
    responseRepairDeliverableFailures: input.responseRepairDeliverableFailures,
    toolStepCount: input.toolStepCount,
    actionPolicyOverride: input.actionPolicy,
  });
  if (input.assertArtifacts) {
    await input.runtime.assertArtifactObligationsSatisfied(input.turn, input.artifactObligations);
  }
  return input.assertTerminal === false
    ? repairedText
    : input.runtime.assertValidTerminalAssistantTextOrFallback(input.turn, repairedText);
}
