import { BestAgentError } from "@best-agent/foundation";
import type {
  ProviderTurnTraceEntry,
  RequestBudgetSnapshot,
  SessionUsageOperation,
  Turn,
} from "@best-agent/protocol";
import type {
  ConversationMessage,
  ProviderCacheControl,
  ProviderToolOption,
} from "../../provider/provider-core.js";
import type { ProviderCorrectionTrace } from "../checkpoint-codec.js";
import { appendRuntimeObservation } from "./runtime-observation.js";
import { validateProviderResponseFinalAnswer } from "./provider-final-answer-handler.js";
import { assertValidProviderProseResponse } from "./provider-response-validation.js";
import type { ProviderTurnInvokeContext } from "./provider-turn-budget.js";

export type GenerateFinalAnswerRepairRuntime = {
  assertValidTerminalAssistantText: (text: string) => string;
  buildCorrectionTrace: (input: {
    code: ProviderCorrectionTrace["code"];
    messages: ConversationMessage[];
    operation: SessionUsageOperation;
    systemPrompt?: string;
    tools?: ProviderToolOption[];
  }) => ProviderCorrectionTrace;
  buildProviderInvokeContext: (
    overrides?: Partial<ProviderTurnInvokeContext>,
  ) => ProviderTurnInvokeContext;
  provider: {
    generate: (
      threadId: Turn["threadId"],
      turnId: Turn["id"],
      request: {
        cacheControl?: ProviderCacheControl;
        messages: ConversationMessage[];
        systemPrompt?: string;
      },
      operation: SessionUsageOperation,
      context: ProviderTurnInvokeContext,
    ) => Promise<{ text: string }>;
  };
};

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

export async function repairDurableGenerateFinalAnswer(input: {
  availableTools: ProviderToolOption[];
  cacheControl?: ProviderCacheControl;
  finalAnswerRepairAttempted: boolean;
  messages: ConversationMessage[];
  objective?: string;
  operation: SessionUsageOperation;
  requestBudgetSnapshot?: RequestBudgetSnapshot;
  responseText: string;
  runtime: GenerateFinalAnswerRepairRuntime;
  systemPrompt?: string;
  turn: Turn;
}): Promise<{
  repairValidation?: ProviderTurnTraceValidation;
  repaired: boolean;
  responseText: string;
}> {
  const finalValidation = validateProviderResponseFinalAnswer({
    availableTools: input.availableTools,
    messages: input.messages,
    objective: input.objective,
    text: input.responseText,
  });
  if (finalValidation.ok) {
    return {
      repaired: false,
      responseText: input.responseText,
    };
  }
  if (input.finalAnswerRepairAttempted) {
    throw new BestAgentError("PROVIDER_DECISION_INVALID", finalValidation.observation.message);
  }
  const repairMessages = appendRuntimeObservation(input.messages, finalValidation.observation);
  const correctionTrace = input.runtime.buildCorrectionTrace({
    code: "PROVIDER_DECISION_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 {
    repairValidation: {
      message: finalValidation.observation.message,
      reasonCode: "final-answer-unsupported",
    },
    repaired: true,
    responseText: assertValidProviderProseResponse({
      assertValidTerminalAssistantText: input.runtime.assertValidTerminalAssistantText,
      availableTools: input.availableTools,
      messages: repairMessages,
      objective: input.objective,
      text: repairResponse.text,
    }),
  };
}
