import type { RequestBudgetSnapshot } from "@best-agent/protocol";
import type {
  ConversationMessage,
  ProviderCacheControl,
  ProviderDecisionOutput,
  ProviderTextStreamCallbacks,
  ProviderToolOption,
} from "../../provider/provider-core.js";
import type {
  ProviderDecisionResolutionControls,
  ProviderDecisionResolutionInput,
} from "./provider-decision-resolution.js";
import type { RequiredToolConstraint } from "./tool-constraint-types.js";

export type ProviderGenerateOnlyPreparedRequest = {
  request: {
    cacheControl?: ProviderCacheControl;
    messages: ConversationMessage[];
    systemPrompt?: string;
  };
  snapshot: RequestBudgetSnapshot;
};

export async function runProviderGenerateOnlyTurn(input: {
  assertTurnAttemptWritable: () => Promise<void>;
  canRetryContextOverflow: (error: unknown, snapshot: RequestBudgetSnapshot) => boolean;
  emitReasoning: (reasoningText: string) => void;
  invokeGenerate: (
    preparedRequest: ProviderGenerateOnlyPreparedRequest,
  ) => Promise<{ reasoningText?: string; text: string }>;
  prepareRequest: (options?: {
    forceCompaction?: boolean;
  }) => Promise<ProviderGenerateOnlyPreparedRequest>;
  repairResponse: (
    preparedRequest: ProviderGenerateOnlyPreparedRequest,
    responseText: string,
  ) => Promise<{
    finalAnswerRepaired: boolean;
    responseText: string;
    terminalOutputRepaired: boolean;
  }>;
  sawReasoning: () => boolean;
}): Promise<{
  finalAnswerRepaired: boolean;
  responseText: string;
  terminalOutputRepaired: boolean;
}> {
  let preparedRequest = await input.prepareRequest();
  await input.assertTurnAttemptWritable();
  let response: Awaited<ReturnType<typeof input.invokeGenerate>>;
  try {
    response = await input.invokeGenerate(preparedRequest);
  } catch (error) {
    if (!input.canRetryContextOverflow(error, preparedRequest.snapshot)) {
      throw error;
    }
    preparedRequest = await input.prepareRequest({ forceCompaction: true });
    await input.assertTurnAttemptWritable();
    response = await input.invokeGenerate(preparedRequest);
  }
  if (!input.sawReasoning() && response.reasoningText) {
    input.emitReasoning(response.reasoningText);
  }
  return input.repairResponse(preparedRequest, response.text);
}

export type ProviderDecisionPreparedRequest = {
  request: {
    cacheControl?: ProviderCacheControl;
    messages: ConversationMessage[];
    systemPrompt?: string;
    tools?: ProviderToolOption[];
  };
  snapshot: RequestBudgetSnapshot;
};

export type ProviderDecisionPreview = {
  callbacks: ProviderTextStreamCallbacks;
  finalize: () => Promise<void>;
  resolveToolCallId: (
    providerToolCallId?: string,
  ) => ProviderDecisionResolutionInput["resolvePreviewToolCallId"] extends
    | ((providerToolCallId?: string) => infer R)
    | undefined
    ? R
    : never;
};

export async function runProviderDecisionTurn(input: {
  approvalTurn?: ProviderDecisionResolutionInput["approvalTurn"];
  assertTurnAttemptWritable: () => Promise<void>;
  availableTools: ProviderToolOption[];
  buildDecideMemoryRequest: (options?: {
    forceCompaction?: boolean;
  }) => Promise<ProviderDecisionPreparedRequest>;
  canRetryContextOverflow: (error: unknown, snapshot: RequestBudgetSnapshot) => boolean;
  controls: ProviderDecisionResolutionControls;
  createDecisionPreview: (
    visibleTools: ProviderToolOption[],
    noActionTimeoutMs?: number,
  ) => ProviderDecisionPreview;
  invokeDecide: (
    preparedRequest: ProviderDecisionPreparedRequest,
    input: {
      decisionMessages: ConversationMessage[];
      decisionPreview: ProviderDecisionPreview;
      decisionSystemPrompt?: string;
      requiredToolConstraint?: RequiredToolConstraint;
      visibleTools: ProviderToolOption[];
    },
  ) => Promise<ProviderDecisionOutput>;
  noActionTimeoutMs: () => number;
  providerTurnBudgeted: boolean;
  resolveDecision: (input: {
    decision: ProviderDecisionOutput;
    decisionMessages: ConversationMessage[];
    decisionPreview: ProviderDecisionPreview;
    decisionSystemPrompt?: string;
    preparedRequest: ProviderDecisionPreparedRequest;
    requiredToolConstraint?: RequiredToolConstraint;
    visibleTools: ProviderToolOption[];
  }) => Promise<string>;
  resolveRequiredToolConstraint: (
    messages: ConversationMessage[],
    availableTools: ProviderToolOption[],
  ) => RequiredToolConstraint | undefined;
  resolveVisibleTools: (
    availableTools: ProviderToolOption[],
    requiredToolConstraint: RequiredToolConstraint | undefined,
    messages: ConversationMessage[],
  ) => ProviderToolOption[];
}): Promise<{
  completionGuardRepaired: boolean;
  responseText: string;
  terminalOutputRepaired: boolean;
}> {
  let preparedRequest = await input.buildDecideMemoryRequest();
  let requiredToolConstraint = input.resolveRequiredToolConstraint(
    preparedRequest.request.messages,
    input.availableTools,
  );
  let providerVisibleTools = input.resolveVisibleTools(
    input.availableTools,
    requiredToolConstraint,
    preparedRequest.request.messages,
  );
  let decision: ProviderDecisionOutput;
  let decisionMessages = preparedRequest.request.messages;
  let decisionSystemPrompt = preparedRequest.request.systemPrompt;
  let decisionPreview = input.createDecisionPreview(
    providerVisibleTools,
    input.providerTurnBudgeted ? input.noActionTimeoutMs() : undefined,
  );
  const invokeCurrentDecision = async (): Promise<ProviderDecisionOutput> => {
    await input.assertTurnAttemptWritable();
    return input.invokeDecide(preparedRequest, {
      decisionMessages,
      decisionPreview,
      decisionSystemPrompt,
      requiredToolConstraint,
      visibleTools: providerVisibleTools,
    });
  };

  try {
    decision = await invokeCurrentDecision();
  } catch (error) {
    if (input.canRetryContextOverflow(error, preparedRequest.snapshot)) {
      await decisionPreview.finalize();
      preparedRequest = await input.buildDecideMemoryRequest({ forceCompaction: true });
      requiredToolConstraint = input.resolveRequiredToolConstraint(
        preparedRequest.request.messages,
        input.availableTools,
      );
      providerVisibleTools = input.resolveVisibleTools(
        input.availableTools,
        requiredToolConstraint,
        preparedRequest.request.messages,
      );
      decisionMessages = preparedRequest.request.messages;
      decisionSystemPrompt = preparedRequest.request.systemPrompt;
      decisionPreview = input.createDecisionPreview(
        providerVisibleTools,
        input.providerTurnBudgeted ? input.noActionTimeoutMs() : undefined,
      );
      decision = await invokeCurrentDecision();
    } else {
      throw error;
    }
  }
  await decisionPreview.finalize();

  const responseText = await input.resolveDecision({
    decision,
    decisionMessages,
    decisionPreview,
    decisionSystemPrompt,
    preparedRequest,
    requiredToolConstraint,
    visibleTools: providerVisibleTools,
  });
  return { completionGuardRepaired: false, responseText, terminalOutputRepaired: false };
}
