type UnknownRecord = Record<string, unknown>;

function asRecord(value: unknown): UnknownRecord | undefined {
  return typeof value === "object" && value !== null && !Array.isArray(value)
    ? value as UnknownRecord
    : undefined;
}

function nonEmptyString(value: unknown): string | undefined {
  return typeof value === "string" && value.length > 0 ? value : undefined;
}

function roleOf(message: unknown): string | undefined {
  return nonEmptyString(asRecord(message)?.role);
}

function assistantCallIds(message: unknown): Set<string> | undefined {
  const assistant = asRecord(message);
  if (assistant?.role !== "assistant" || !Array.isArray(assistant.content)) return undefined;

  const ids = new Set<string>();
  for (const content of assistant.content) {
    const block = asRecord(content);
    if (block?.type !== "toolCall") continue;

    const id = nonEmptyString(block.id);
    const name = nonEmptyString(block.name);
    if (!id || !name || !asRecord(block.arguments) || ids.has(id)) return undefined;
    ids.add(id);
  }

  return ids.size > 0 ? ids : undefined;
}

function resultCallId(message: unknown): string | undefined {
  const result = asRecord(message);
  if (result?.role !== "toolResult") return undefined;
  if (!nonEmptyString(result.toolName)) return undefined;
  if (!Array.isArray(result.content) || typeof result.isError !== "boolean") return undefined;
  return nonEmptyString(result.toolCallId);
}

interface TrailingBatch {
  assistant: unknown;
  resultIds: Set<string>;
}

function trailingBatch(messages: readonly unknown[]): TrailingBatch | undefined {
  if (roleOf(messages.at(-1)) !== "toolResult") return undefined;

  const resultIds = new Set<string>();
  let index = messages.length - 1;
  while (index >= 0 && roleOf(messages[index]) === "toolResult") {
    const id = resultCallId(messages[index]);
    if (!id || resultIds.has(id)) return undefined;
    resultIds.add(id);
    index -= 1;
  }

  if (index < 0 || roleOf(messages[index]) !== "assistant") return undefined;
  return { assistant: messages[index], resultIds };
}

/** True when the context ends with one complete Pi assistant/tool-result batch. */
export function endsWithCompleteToolResultBatch(messages: readonly unknown[]): boolean {
  const batch = trailingBatch(messages);
  if (!batch) return false;

  const callIds = assistantCallIds(batch.assistant);
  if (!callIds || callIds.size !== batch.resultIds.size) return false;

  for (const id of callIds) {
    if (!batch.resultIds.has(id)) return false;
  }
  return true;
}