import type { SessionUsageOperation, Turn } from "@best-agent/protocol";
import type { ConversationMessage, ProviderToolOption } from "../../provider/model-provider.js";
import type {
  ProviderCorrectionTrace,
  ProviderGenerateCheckpointRequest,
} from "../checkpoint-codec.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";

export async function runReplayGenerateRepair(input: {
  availableTools: ProviderToolOption[];
  baseMessages: ConversationMessage[];
  correctionOperation: SessionUsageOperation;
  objective?: string;
  providerContext: ProviderTurnInvokeContext;
  repairCode: ProviderCorrectionTrace["code"];
  repairObservation: RuntimeObservation;
  request: ProviderGenerateCheckpointRequest;
  runtime: Pick<
    ProviderReplayRepairRuntime,
    "buildCorrectionTrace" | "finalizeGeneratedResponse" | "provider"
  >;
  systemPrompt?: string;
  traceSystemPrompt?: string;
  turn: Turn;
}): Promise<string> {
  const repairMessages = appendRuntimeObservation(input.baseMessages, input.repairObservation);
  const correctionTrace = input.runtime.buildCorrectionTrace({
    code: input.repairCode,
    messages: input.baseMessages,
    operation: input.correctionOperation,
    systemPrompt: input.traceSystemPrompt ?? input.systemPrompt,
  });
  const response = await input.runtime.provider.generate(
    input.turn.threadId,
    input.turn.id,
    {
      ...input.request,
      messages: repairMessages,
      systemPrompt: input.systemPrompt,
    },
    "repair",
    {
      ...input.providerContext,
      correctionTrace,
    },
  );
  return input.runtime.finalizeGeneratedResponse({
    availableTools: input.availableTools,
    messages: repairMessages,
    objective: input.objective,
    text: response.text,
  });
}
