import type {
  RequestBudgetSnapshot,
  SessionUsageOperation,
  Turn,
  TurnSettlementRecord,
} from "@best-agent/protocol";
import type {
  ConversationMessage,
  ProviderCacheControl,
  ProviderToolOption,
} from "../../provider/provider-core.js";
import type { PromptDeliverableFailure } from "../../prompt/runtime-prompts.js";
import { classifyRuntimeFailure } from "../failure-classification.js";
import type { ProviderCorrectionTrace } from "../checkpoint-codec.js";
import type { ProviderTurnInvokeContext } from "./provider-turn-budget.js";
import { buildTerminalOutputInvalidRepair } from "./provider-terminal-output-handler.js";
import { appendRuntimeObservation } from "./runtime-observation.js";

type GenerateTerminalRepairRuntime = {
  assertValidTerminalAssistantText: (text: string) => string;
  buildCorrectionTrace: (input: {
    code: ProviderCorrectionTrace["code"];
    messages: ConversationMessage[];
    operation: SessionUsageOperation;
    systemPrompt?: string;
    tools?: ProviderToolOption[];
  }) => ProviderCorrectionTrace;
  buildProviderInvokeContext: (
    overrides?: Partial<ProviderTurnInvokeContext>,
  ) => ProviderTurnInvokeContext;
  buildTurnSettlementRecord: (input: {
    rejectedReason?: TurnSettlementRecord["rejectedReason"];
    terminalKind: TurnSettlementRecord["terminalKind"];
    unresolvedFailures?: PromptDeliverableFailure[];
  }) => Promise<TurnSettlementRecord>;
  provider: {
    generate: (
      threadId: Turn["threadId"],
      turnId: Turn["id"],
      request: {
        cacheControl?: ProviderCacheControl;
        messages: ConversationMessage[];
        systemPrompt?: string;
      },
      operation: SessionUsageOperation,
      context: ProviderTurnInvokeContext,
    ) => Promise<{ text: string }>;
  };
  recordTurnSettlement: (settlement: TurnSettlementRecord) => Promise<void>;
};

export async function repairDurableGenerateTerminalOutput(input: {
  availableTools: ProviderToolOption[];
  cacheControl?: ProviderCacheControl;
  messages: ConversationMessage[];
  operation: SessionUsageOperation;
  repairAttempted: boolean;
  requestBudgetSnapshot?: RequestBudgetSnapshot;
  responseText: string;
  runtime: GenerateTerminalRepairRuntime;
  systemPrompt?: string;
  turn: Turn;
}): Promise<{ repaired: boolean; responseText: string }> {
  try {
    return {
      repaired: false,
      responseText: input.runtime.assertValidTerminalAssistantText(input.responseText),
    };
  } catch (error) {
    const failure = classifyRuntimeFailure(error);
    if (failure.errorCode !== "TERMINAL_OUTPUT_INVALID" || input.repairAttempted) {
      throw error;
    }
  }

  await input.runtime.recordTurnSettlement(
    await input.runtime.buildTurnSettlementRecord({
      rejectedReason: "invalid-terminal-output",
      terminalKind: "completed",
      unresolvedFailures: [],
    }),
  );
  const terminalRepair = buildTerminalOutputInvalidRepair({
    baseActionPolicy: {
      allowedActions: ["respond"],
      directResponseAllowed: true,
    },
    classificationToolNames: input.availableTools.map((tool) => tool.name),
    deliverableFailures: [],
    text: input.responseText,
    visibleTools: [],
  });
  const repairMessages = appendRuntimeObservation(input.messages, terminalRepair.observation);
  const correctionTrace = input.runtime.buildCorrectionTrace({
    code: "TERMINAL_OUTPUT_INVALID",
    messages: input.messages,
    operation: input.operation,
    systemPrompt: input.systemPrompt,
  });
  const repairResponse = await input.runtime.provider.generate(
    input.turn.threadId,
    input.turn.id,
    {
      cacheControl: input.cacheControl,
      messages: repairMessages,
      systemPrompt: input.systemPrompt,
    },
    "repair",
    input.runtime.buildProviderInvokeContext({
      correctionTrace,
      requestBudgetSnapshot: input.requestBudgetSnapshot,
      supersedePendingProviderRequest: true,
    }),
  );
  return {
    repaired: true,
    responseText: input.runtime.assertValidTerminalAssistantText(repairResponse.text),
  };
}
