import type {
  ProviderTurnTraceEntry,
  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 type { ProviderCorrectionTrace } from "../checkpoint-codec.js";
import type { ProviderTurnInvokeContext } from "./provider-turn-budget.js";
import { repairDurableGenerateTerminalOutput } from "./provider-terminal-repair-runner.js";
import { repairDurableGenerateFinalAnswer } from "./provider-final-answer-repair-runner.js";
import { validateProviderResponseFinalAnswer } from "./provider-final-answer-handler.js";

export { repairDurableGenerateTerminalOutput } from "./provider-terminal-repair-runner.js";
export { repairDurableGenerateFinalAnswer } from "./provider-final-answer-repair-runner.js";

type ProviderTurnTraceValidation = NonNullable<ProviderTurnTraceEntry["validation"]>;

export type RepairGenerateTraceOutcome = {
  runtimeOutcome: "accepted" | "rejected" | "repaired";
  validation?: ProviderTurnTraceValidation;
};

const repairGenerateTraceOutcomesKey = Symbol("repairGenerateTraceOutcomes");

export function getRepairGenerateTraceOutcomes(
  error: unknown,
): RepairGenerateTraceOutcome[] | undefined {
  return typeof error === "object" && error !== null
    ? (error as { [repairGenerateTraceOutcomesKey]?: RepairGenerateTraceOutcome[] })[
        repairGenerateTraceOutcomesKey
      ]
    : undefined;
}

function throwWithRepairGenerateTraceOutcomes(
  error: unknown,
  outcomes: RepairGenerateTraceOutcome[],
): never {
  if (typeof error === "object" && error !== null && outcomes.length > 0) {
    Object.defineProperty(error, repairGenerateTraceOutcomesKey, {
      configurable: true,
      value: outcomes,
    });
  }
  throw error;
}

export type DurableResponseRepairRuntime = {
  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: ProviderResponseRepairProvider;
  recordTurnSettlement: (settlement: TurnSettlementRecord) => Promise<void>;
};

type ProviderResponseRepairProvider = {
  generate: (
    threadId: Turn["threadId"],
    turnId: Turn["id"],
    request: {
      cacheControl?: ProviderCacheControl;
      messages: ConversationMessage[];
      systemPrompt?: string;
    },
    operation: SessionUsageOperation,
    context: ProviderTurnInvokeContext,
  ) => Promise<{ text: string }>;
};

export async function repairDurableGenerateResponse(input: {
  availableTools: ProviderToolOption[];
  cacheControl?: ProviderCacheControl;
  finalAnswerRepairAttempted: boolean;
  messages: ConversationMessage[];
  objective?: string;
  operation: SessionUsageOperation;
  repairAttempted: boolean;
  requestBudgetSnapshot?: RequestBudgetSnapshot;
  responseText: string;
  runtime: DurableResponseRepairRuntime;
  systemPrompt?: string;
  turn: Turn;
  validateFinalAnswer?: boolean;
}): Promise<{
  finalAnswerRepaired: boolean;
  repairGenerateOutcomes: RepairGenerateTraceOutcome[];
  responseText: string;
  terminalOutputRepaired: boolean;
}> {
  const terminalRepair = await repairDurableGenerateTerminalOutput(input);
  const repairGenerateOutcomes: RepairGenerateTraceOutcome[] = [];
  const responseText = terminalRepair.responseText;
  if (input.validateFinalAnswer === false) {
    if (terminalRepair.repaired) {
      repairGenerateOutcomes.push({ runtimeOutcome: "accepted" });
    }
    return {
      finalAnswerRepaired: false,
      repairGenerateOutcomes,
      responseText,
      terminalOutputRepaired: terminalRepair.repaired,
    };
  }
  const terminalRepairFinalValidation = terminalRepair.repaired
    ? validateProviderResponseFinalAnswer({
        availableTools: input.availableTools,
        messages: input.messages,
        objective: input.objective,
        text: responseText,
      })
    : undefined;
  let finalAnswerRepair: Awaited<ReturnType<typeof repairDurableGenerateFinalAnswer>>;
  try {
    finalAnswerRepair = await repairDurableGenerateFinalAnswer({
      availableTools: input.availableTools,
      cacheControl: input.cacheControl,
      finalAnswerRepairAttempted: input.finalAnswerRepairAttempted,
      messages: input.messages,
      objective: input.objective,
      operation: input.operation,
      requestBudgetSnapshot: input.requestBudgetSnapshot,
      responseText,
      runtime: input.runtime,
      systemPrompt: input.systemPrompt,
      turn: input.turn,
    });
  } catch (error) {
    if (terminalRepair.repaired) {
      throwWithRepairGenerateTraceOutcomes(error, [
        {
          runtimeOutcome: "repaired",
          validation:
            terminalRepairFinalValidation && !terminalRepairFinalValidation.ok
              ? {
                  message: terminalRepairFinalValidation.observation.message,
                  reasonCode: "final-answer-unsupported",
                }
              : {
                  message: "Generate repair response was superseded by final-answer repair.",
                  reasonCode: "final-answer-unsupported",
                },
        },
      ]);
    }
    throw error;
  }
  if (terminalRepair.repaired) {
    repairGenerateOutcomes.push(
      finalAnswerRepair.repaired
        ? {
            runtimeOutcome: "repaired",
            validation: finalAnswerRepair.repairValidation ?? {
              message: "Generate repair response was superseded by final-answer repair.",
              reasonCode: "final-answer-unsupported",
            },
          }
        : { runtimeOutcome: "accepted" },
    );
  }
  if (finalAnswerRepair.repaired) {
    repairGenerateOutcomes.push({ runtimeOutcome: "accepted" });
  }
  return {
    finalAnswerRepaired: finalAnswerRepair.repaired,
    repairGenerateOutcomes,
    responseText: finalAnswerRepair.responseText,
    terminalOutputRepaired: terminalRepair.repaired,
  };
}
