import type { CompletedProviderToolCall } from "../checkpoint-codec.js";
import {
  requiredAnyToolGroupsForConstraint,
  splitRequiredToolGroups,
  type RequiredToolConstraint,
} from "./tool-constraint-types.js";

export function normalizeRequiredToolConstraint(
  input: RequiredToolConstraint,
): RequiredToolConstraint | undefined {
  const forbiddenToolNames = uniqueToolNames(input.forbiddenToolNames);
  const forbidden = new Set(forbiddenToolNames);
  const originalRequiredToolNames = uniqueToolNames(input.requiredToolNames);
  const originalRequiredAnyToolGroups = [
    ...(input.requiredAnyToolGroups ?? []),
    ...(input.requiredAnyToolNames ? [input.requiredAnyToolNames] : []),
  ];
  const hadOriginalToolRequirement =
    originalRequiredToolNames.length > 0 ||
    (input.requiredAllToolNames?.length ?? 0) > 0 ||
    originalRequiredAnyToolGroups.some((group) => group.length > 0);
  const requiredAnyToolGroupsAfterForbidden = originalRequiredAnyToolGroups
    .map((group) => uniqueToolNames(group).filter((toolName) => !forbidden.has(toolName)))
    .filter((group) => group.length > 0);
  const normalizedRequiredAnyToolGroups = requiredAnyToolGroupsAfterForbidden.filter(
    (group) => group.length > 1,
  );
  const requiredAnyToolNames = new Set(normalizedRequiredAnyToolGroups.flat());
  const originalRequiredAllToolNames = input.requiredAllToolNames
    ? uniqueToolNames(input.requiredAllToolNames)
    : originalRequiredToolNames.filter((toolName) => !requiredAnyToolNames.has(toolName));
  const collapsedRequiredAnyToolNames = requiredAnyToolGroupsAfterForbidden
    .filter((group) => group.length === 1)
    .map((group) => group[0]!);
  const normalizedRequiredAllToolNames = uniqueToolNames([
    ...originalRequiredAllToolNames.filter((toolName) => !forbidden.has(toolName)),
    ...collapsedRequiredAnyToolNames,
  ]);
  const normalizedRequiredAll = new Set(normalizedRequiredAllToolNames);
  const normalizedRequiredToolNames: string[] = [];
  const addNormalizedRequiredToolName = (toolName: string) => {
    if (!normalizedRequiredToolNames.includes(toolName)) {
      normalizedRequiredToolNames.push(toolName);
    }
  };
  for (const toolName of originalRequiredToolNames) {
    if (forbidden.has(toolName)) {
      continue;
    }
    if (
      normalizedRequiredAll.has(toolName) ||
      normalizedRequiredAnyToolGroups.some((group) => group.includes(toolName))
    ) {
      addNormalizedRequiredToolName(toolName);
    }
  }
  for (const toolName of normalizedRequiredAllToolNames) {
    addNormalizedRequiredToolName(toolName);
  }
  for (const group of normalizedRequiredAnyToolGroups) {
    if (!normalizedRequiredToolNames.some((toolName) => group.includes(toolName))) {
      addNormalizedRequiredToolName(group[0]!);
    }
  }
  const normalizedRequiredAnyToolNames =
    normalizedRequiredAnyToolGroups.length === 1 ? normalizedRequiredAnyToolGroups[0] : undefined;
  const hasSelectableToolNames = input.selectableToolNames !== undefined;
  const defaultSelectableToolNames = uniqueToolNames([
    ...normalizedRequiredToolNames,
    ...normalizedRequiredAnyToolGroups.flat(),
  ]);
  const selectableToolNames = uniqueToolNames(
    hasSelectableToolNames ? (input.selectableToolNames ?? []) : defaultSelectableToolNames,
  ).filter((toolName) => !forbidden.has(toolName));
  const visibleToolNames = input.visibleToolNames
    ? uniqueToolNames(input.visibleToolNames).filter((toolName) => !forbidden.has(toolName))
    : undefined;
  const selectable = new Set(selectableToolNames);
  const visible = visibleToolNames ? new Set(visibleToolNames) : undefined;
  const requiredDroppedByPolicy = normalizedRequiredToolNames.filter(
    (toolName) => !selectable.has(toolName) || (visible && !visible.has(toolName)),
  );
  const requiredAnyDroppedByPolicy = normalizedRequiredAnyToolGroups
    .flat()
    .filter((toolName) => !selectable.has(toolName) || (visible && !visible.has(toolName)));

  if (requiredDroppedByPolicy.length > 0 || requiredAnyDroppedByPolicy.length > 0) {
    throw new Error(
      `Required tools are not selectable/visible: ${[
        ...requiredDroppedByPolicy,
        ...requiredAnyDroppedByPolicy,
      ].join(", ")}`,
    );
  }

  if (
    input.requiredReleasePolicy === "after-required-tools-succeeded" &&
    normalizedRequiredToolNames.length === 0
  ) {
    if (forbiddenToolNames.length > 0 && hadOriginalToolRequirement) {
      const forbiddenOnlyConstraint: RequiredToolConstraint = {
        ...(input.conflicts ? { conflicts: input.conflicts } : {}),
        ...(input.evidence ? { evidence: input.evidence } : {}),
        forbiddenToolNames,
        requiredReleasePolicy: "none",
        requiredToolNames: [],
        reason: "explicit-forbidden-tool",
        restrictionScope: "turn",
        ...(visibleToolNames && visibleToolNames.length > 0 ? { visibleToolNames } : {}),
      };
      assertRequiredToolConstraintInvariants(forbiddenOnlyConstraint);
      return forbiddenOnlyConstraint;
    }
    throw new Error("Required tool constraint has release policy but no required tools.");
  }

  const primaryToolName =
    normalizedRequiredToolNames[0] ??
    (input.primaryToolName && !forbidden.has(input.primaryToolName)
      ? input.primaryToolName
      : selectableToolNames[0]);
  const shouldPreserveVisibleToolNames =
    visibleToolNames !== undefined && visibleToolNames.length > 0;
  const shouldPreserveSelectableToolNames =
    selectableToolNames.length > 0 &&
    (hasSelectableToolNames || normalizedRequiredToolNames.length > 0);
  const normalized: RequiredToolConstraint = {
    ...(input.conflicts ? { conflicts: input.conflicts } : {}),
    ...(input.evidence ? { evidence: input.evidence } : {}),
    forbiddenToolNames,
    ...(primaryToolName ? { primaryToolName } : {}),
    reason: input.reason,
    ...(normalizedRequiredAllToolNames.length > 0
      ? { requiredAllToolNames: normalizedRequiredAllToolNames }
      : {}),
    ...(normalizedRequiredAnyToolNames
      ? { requiredAnyToolNames: normalizedRequiredAnyToolNames }
      : {}),
    ...(normalizedRequiredAnyToolGroups.length > 1
      ? { requiredAnyToolGroups: normalizedRequiredAnyToolGroups }
      : {}),
    requiredReleasePolicy: input.requiredReleasePolicy,
    requiredToolNames: normalizedRequiredToolNames,
    restrictionScope: input.restrictionScope,
    ...(shouldPreserveSelectableToolNames ? { selectableToolNames } : {}),
    ...(shouldPreserveVisibleToolNames ? { visibleToolNames } : {}),
  };
  assertRequiredToolConstraintInvariants(normalized);
  return normalized;
}

function assertRequiredToolConstraintInvariants(constraint: RequiredToolConstraint): void {
  const forbidden = new Set(constraint.forbiddenToolNames);
  const selectable = constraint.selectableToolNames
    ? new Set(constraint.selectableToolNames)
    : undefined;
  const visible = constraint.visibleToolNames ? new Set(constraint.visibleToolNames) : undefined;

  if (
    constraint.requiredReleasePolicy === "none" &&
    (constraint.requiredToolNames.length > 0 ||
      (constraint.requiredAllToolNames?.length ?? 0) > 0 ||
      requiredAnyToolGroupsForConstraint(constraint).length > 0)
  ) {
    throw new Error("requiredReleasePolicy=none cannot have requiredToolNames.");
  }

  if (
    constraint.requiredReleasePolicy === "after-required-tools-succeeded" &&
    constraint.requiredToolNames.length === 0
  ) {
    throw new Error("Required release policy needs at least one required tool.");
  }

  for (const toolName of constraint.forbiddenToolNames) {
    if (selectable?.has(toolName)) {
      throw new Error(`Forbidden tool ${toolName} is selectable.`);
    }
    if (visible?.has(toolName)) {
      throw new Error(`Forbidden tool ${toolName} is visible.`);
    }
  }

  if (visible && selectable) {
    for (const toolName of selectable) {
      if (!visible.has(toolName)) {
        throw new Error(`Selectable tool ${toolName} is not visible.`);
      }
    }
  }

  for (const toolName of constraint.requiredToolNames) {
    if (forbidden.has(toolName)) {
      throw new Error(`Required tool ${toolName} is also forbidden.`);
    }
    if (selectable && !selectable.has(toolName)) {
      throw new Error(`Required tool ${toolName} is not selectable.`);
    }
    if (visible && !visible.has(toolName)) {
      throw new Error(`Required tool ${toolName} is not visible.`);
    }
  }

  for (const toolName of constraint.requiredAllToolNames ?? []) {
    if (!constraint.requiredToolNames.includes(toolName)) {
      throw new Error(`All-required tool ${toolName} is missing from requiredToolNames.`);
    }
    if (forbidden.has(toolName)) {
      throw new Error(`Required tool ${toolName} is also forbidden.`);
    }
    if (selectable && !selectable.has(toolName)) {
      throw new Error(`Required tool ${toolName} is not selectable.`);
    }
    if (visible && !visible.has(toolName)) {
      throw new Error(`Required tool ${toolName} is not visible.`);
    }
  }

  for (const toolName of requiredAnyToolGroupsForConstraint(constraint).flat()) {
    if (forbidden.has(toolName)) {
      throw new Error(`Required tool ${toolName} is also forbidden.`);
    }
    if (selectable && !selectable.has(toolName)) {
      throw new Error(`Required tool ${toolName} is not selectable.`);
    }
    if (visible && !visible.has(toolName)) {
      throw new Error(`Required tool ${toolName} is not visible.`);
    }
  }

  if (constraint.primaryToolName) {
    if (forbidden.has(constraint.primaryToolName)) {
      throw new Error(`Primary tool ${constraint.primaryToolName} is forbidden.`);
    }
    if (selectable && !selectable.has(constraint.primaryToolName)) {
      throw new Error(`Primary tool ${constraint.primaryToolName} is not selectable.`);
    }
    if (visible && !visible.has(constraint.primaryToolName)) {
      throw new Error(`Primary tool ${constraint.primaryToolName} is not visible.`);
    }
  }
}

export function resolveUnresolvedToolConstraint(
  constraint: RequiredToolConstraint | undefined,
  completedToolCalls: CompletedProviderToolCall[],
): RequiredToolConstraint | undefined {
  if (!constraint) {
    return undefined;
  }

  const completedSuccessfulToolNames = new Set(
    completedToolCalls
      .filter((toolCall) => toolCall.outcome !== "error")
      .map((toolCall) => toolCall.toolName),
  );
  const { requiredAllToolNames, requiredAnyToolGroups } = splitRequiredToolGroups(constraint);
  const isRequiredAnyGroupSatisfied = (group: readonly string[]) =>
    group.some((toolName) => completedSuccessfulToolNames.has(toolName));
  const unsatisfiedRequiredAnyToolGroups = requiredAnyToolGroups.filter(
    (group) => !isRequiredAnyGroupSatisfied(group),
  );
  const unsatisfiedRequiredAnyGroupIndexes = new Set(
    requiredAnyToolGroups
      .map((group, index) => (isRequiredAnyGroupSatisfied(group) ? -1 : index))
      .filter((index) => index >= 0),
  );
  const remainingRequiredAllToolNames = requiredAllToolNames.filter(
    (toolName) => !completedSuccessfulToolNames.has(toolName),
  );
  const remainingRequiredAll = new Set(remainingRequiredAllToolNames);
  const remainingRequiredToolNames: string[] = [];
  const addedAnyGroups = new Set<number>();
  const addRemainingRequiredToolName = (toolName: string) => {
    if (!remainingRequiredToolNames.includes(toolName)) {
      remainingRequiredToolNames.push(toolName);
    }
  };
  for (const toolName of constraint.requiredToolNames) {
    if (remainingRequiredAll.has(toolName)) {
      addRemainingRequiredToolName(toolName);
      continue;
    }
    const anyGroupIndex = requiredAnyToolGroups.findIndex((group) => group.includes(toolName));
    if (anyGroupIndex >= 0) {
      if (unsatisfiedRequiredAnyGroupIndexes.has(anyGroupIndex)) {
        const group = requiredAnyToolGroups[anyGroupIndex]!;
        if (!addedAnyGroups.has(anyGroupIndex)) {
          addRemainingRequiredToolName(
            group.find((candidate) => !completedSuccessfulToolNames.has(candidate)) ?? group[0]!,
          );
          addedAnyGroups.add(anyGroupIndex);
        }
      }
      continue;
    }
    if (!completedSuccessfulToolNames.has(toolName)) {
      addRemainingRequiredToolName(toolName);
    }
  }
  for (const toolName of remainingRequiredAllToolNames) {
    addRemainingRequiredToolName(toolName);
  }
  for (let index = 0; index < requiredAnyToolGroups.length; index += 1) {
    if (addedAnyGroups.has(index)) {
      continue;
    }
    const group = requiredAnyToolGroups[index]!;
    if (isRequiredAnyGroupSatisfied(group)) {
      continue;
    }
    if (!remainingRequiredToolNames.some((toolName) => group.includes(toolName))) {
      addRemainingRequiredToolName(
        group.find((toolName) => !completedSuccessfulToolNames.has(toolName)) ?? group[0]!,
      );
    }
  }

  if (remainingRequiredToolNames.length === 0 && unsatisfiedRequiredAnyToolGroups.length === 0) {
    if (constraint.reason === "explicit-exclusive-tool-set") {
      return normalizeRequiredToolConstraint({
        ...(constraint.evidence ? { evidence: constraint.evidence } : {}),
        forbiddenToolNames: constraint.forbiddenToolNames,
        ...(constraint.primaryToolName ? { primaryToolName: constraint.primaryToolName } : {}),
        reason: constraint.reason,
        requiredReleasePolicy: "none",
        requiredToolNames: [],
        restrictionScope: constraint.restrictionScope,
        ...(constraint.selectableToolNames
          ? { selectableToolNames: constraint.selectableToolNames }
          : {}),
        ...(constraint.visibleToolNames ? { visibleToolNames: constraint.visibleToolNames } : {}),
      });
    }

    if (constraint.forbiddenToolNames.length > 0) {
      return normalizeRequiredToolConstraint({
        ...(constraint.evidence ? { evidence: constraint.evidence } : {}),
        forbiddenToolNames: constraint.forbiddenToolNames,
        requiredReleasePolicy: "none",
        requiredToolNames: [],
        reason: "explicit-forbidden-tool",
        restrictionScope: "turn",
      });
    }

    return undefined;
  }

  const nextConstraint: RequiredToolConstraint = {
    ...constraint,
    primaryToolName: remainingRequiredToolNames[0],
    ...(unsatisfiedRequiredAnyToolGroups.length === 1
      ? { requiredAnyToolNames: unsatisfiedRequiredAnyToolGroups[0] }
      : {}),
    ...(unsatisfiedRequiredAnyToolGroups.length > 1
      ? { requiredAnyToolGroups: unsatisfiedRequiredAnyToolGroups }
      : {}),
    requiredAllToolNames: remainingRequiredAllToolNames,
    requiredReleasePolicy: "after-required-tools-succeeded",
    requiredToolNames: remainingRequiredToolNames,
  };
  if (unsatisfiedRequiredAnyToolGroups.length !== 1) {
    delete nextConstraint.requiredAnyToolNames;
  }
  if (unsatisfiedRequiredAnyToolGroups.length <= 1) {
    delete nextConstraint.requiredAnyToolGroups;
  }
  if (constraint.reason === "explicit-required-tool") {
    nextConstraint.selectableToolNames =
      unsatisfiedRequiredAnyToolGroups.length > 0
        ? uniqueToolNames([
            ...unsatisfiedRequiredAnyToolGroups.flat(),
            ...remainingRequiredAllToolNames,
          ])
        : remainingRequiredToolNames;
  }
  return normalizeRequiredToolConstraint(nextConstraint);
}

function uniqueToolNames(toolNames: string[]): string[] {
  return [...new Set(toolNames)];
}
