import { BestAgentError } from "@best-agent/foundation";
import type {
  AssetId,
  ToolInputValue,
  ToolReplayPolicy,
} from "@best-agent/protocol";
import type {
  PendingProviderToolCall,
  PersistedToolResumeContext,
} from "../checkpoint-codec.js";
import { isToolInputValue } from "../../tool/tool-history-input.js";

export async function resolveReplayToolInput(
  call: PendingProviderToolCall,
  resolveAsset: (assetId: AssetId, expectedDigest: string | undefined) => Promise<ToolInputValue>,
): Promise<ToolInputValue> {
  if (!call.replayToolInputAssetId) {
    return call.toolInput;
  }
  return resolveAsset(call.replayToolInputAssetId, call.replayToolInputDigest);
}

export async function resolveActiveReplayToolInput(
  payload: PersistedToolResumeContext,
  resolveAsset: (assetId: AssetId, expectedDigest: string | undefined) => Promise<ToolInputValue>,
): Promise<ToolInputValue> {
  if (!payload.activeReplayToolInputAssetId) {
    return payload.toolInput;
  }
  return resolveAsset(payload.activeReplayToolInputAssetId, payload.activeReplayToolInputDigest);
}

export async function resolveReplayToolInputAsset(input: {
  assetId: AssetId;
  expectedDigest?: string;
  readAssetData: (assetId: AssetId) => Promise<Buffer | undefined>;
  toolInputDigest: (value: ToolInputValue) => string;
}): Promise<ToolInputValue> {
  const data = await input.readAssetData(input.assetId);
  if (!data) {
    throw new BestAgentError(
      "RUN_CHECKPOINT_INVALID",
      `Checkpoint replay input asset ${input.assetId} is missing.`,
    );
  }
  let parsed: unknown;
  try {
    parsed = JSON.parse(data.toString("utf8"));
  } catch {
    throw new BestAgentError(
      "RUN_CHECKPOINT_INVALID",
      `Checkpoint replay input asset ${input.assetId} is not valid JSON.`,
    );
  }
  if (!isToolInputValue(parsed)) {
    throw new BestAgentError(
      "RUN_CHECKPOINT_INVALID",
      `Checkpoint replay input asset ${input.assetId} is not a tool input value.`,
    );
  }
  const digest = input.toolInputDigest(parsed);
  if (input.expectedDigest && input.expectedDigest !== digest) {
    throw new BestAgentError(
      "RUN_CHECKPOINT_INVALID",
      `Checkpoint replay input asset ${input.assetId} digest mismatch.`,
    );
  }
  return parsed;
}

export function toolReplayPolicyForFollowUp(input: {
  excludedToolNames: readonly string[];
  resolveReplayPolicy: (toolName: string) => ToolReplayPolicy | undefined;
  toolName: string;
}): ToolReplayPolicy | undefined {
  if (input.excludedToolNames.includes(input.toolName)) {
    return undefined;
  }
  return input.resolveReplayPolicy(input.toolName);
}
