import type { ProviderConfig, ProviderTurnActionPolicy } from "@best-agent/protocol";
import type { AgentProfile } from "../../agent/profile-registry.js";
import {
  mergeSystemPrompts,
  type ConversationMessage,
  type ProviderToolOption,
} from "../../provider/model-provider.js";
import {
  buildToolFollowUpSystemPrompt,
  stripRequiredToolSystemPrompt,
} from "../../prompt/runtime-prompts.js";
import type {
  CompletedProviderToolCall,
  ProviderCorrectionTrace,
  ToolResumeContext,
} from "../checkpoint-codec.js";
import { isRoutingDowngradeEligible } from "./provider-replay-context.js";
import type { ProviderTurnBudget, ProviderTurnInvokeContext } from "./provider-turn-budget.js";
import { appendRuntimeObservation, type RuntimeObservation } from "./runtime-observation.js";
import { resolveProviderVisibleTools } from "./tool-constraint.js";
import type { RequiredToolConstraint } from "./tool-constraint-types.js";

export function buildReplayToolOutputRepairContexts(input: {
  completedToolCalls: CompletedProviderToolCall[];
  fixedReplayRoute: boolean;
  payload: ToolResumeContext;
  repairRequiredToolConstraint?: RequiredToolConstraint;
  repairRouteConfig: ProviderConfig;
  replayProfile: AgentProfile;
  replayProviderTurnBudget: ProviderTurnBudget;
}): {
  providerContext: ProviderTurnInvokeContext;
  resolveContext: ProviderTurnInvokeContext;
  toolStepCount: number;
} {
  const toolStepCount = input.payload.toolStepCount + input.completedToolCalls.length;
  const base = {
    activeSkillIds: input.payload.activeSkillIds,
    advancedPhase: input.payload.advancedPhase,
    artifactObligations: input.payload.artifactObligations,
    availableTools: input.payload.availableTools,
    budget: input.replayProviderTurnBudget,
    downgradeEligible: isRoutingDowngradeEligible(input.payload.phaseName),
    objective: input.payload.objective,
    phaseName: input.payload.phaseName,
    profile: input.replayProfile,
    profileName: input.payload.profileName,
    requiredToolConstraint: input.repairRequiredToolConstraint,
    routeConfigOverride: input.repairRouteConfig,
    toolStepCount,
  };
  return {
    providerContext: {
      ...base,
      fixedRouteConfigOverride: input.fixedReplayRoute,
    },
    resolveContext: {
      ...base,
      fixedRouteConfigOverride: input.fixedReplayRoute,
    },
    toolStepCount,
  };
}

export function prepareReplayToolOutputRepair(input: {
  completedToolCalls: CompletedProviderToolCall[];
  payload: ToolResumeContext;
  repairActionPolicyOverride?: ProviderTurnActionPolicy;
  repairCode: ProviderCorrectionTrace["code"];
  repairMessages: ConversationMessage[];
  repairObservation: RuntimeObservation;
  requiresDeliverableAssertion?: boolean;
  suppressStructuredControls?: boolean;
  unresolvedRequiredToolConstraint?: RequiredToolConstraint;
}): {
  providerRepairMessages: ConversationMessage[];
  repairActionPolicy?: ProviderTurnActionPolicy;
  repairRequiredToolConstraint?: RequiredToolConstraint;
  repairSystemPrompt?: string;
  requiresDeliverableAssertion: boolean;
  visibleTools: ProviderToolOption[];
} {
  const requiresDeliverableAssertion =
    input.requiresDeliverableAssertion || input.repairCode === "PROMPT_DELIVERABLE_MISSING";
  const suppressStructuredControls =
    input.suppressStructuredControls || requiresDeliverableAssertion;
  const providerRepairMessages = appendRuntimeObservation(
    input.repairMessages,
    input.repairObservation,
  );
  const repairRequiredToolConstraint = suppressStructuredControls
    ? undefined
    : input.unresolvedRequiredToolConstraint;
  return {
    providerRepairMessages,
    repairActionPolicy:
      input.repairActionPolicyOverride ??
      (suppressStructuredControls
        ? {
            allowedActions: ["tool"],
            directResponseAllowed: false,
          }
        : undefined),
    repairRequiredToolConstraint,
    repairSystemPrompt: mergeSystemPrompts(
      stripRequiredToolSystemPrompt(
        input.payload.systemPrompt,
        input.payload.requiredToolConstraint,
      ),
      buildToolFollowUpSystemPrompt(input.completedToolCalls),
    ),
    requiresDeliverableAssertion,
    visibleTools: resolveProviderVisibleTools(
      input.payload.availableTools,
      repairRequiredToolConstraint,
      providerRepairMessages,
    ),
  };
}
