import type { WorkerCatalog } from "../catalog/definition.ts";

const CONTRACT_START = "<!-- pi-orchestrate:contract:start -->";
const CONTRACT_END = "<!-- pi-orchestrate:contract:end -->";

function sortedWorkers(catalog: WorkerCatalog) {
  return [...catalog.workers].sort((left, right) => {
    if (left.name < right.name) return -1;
    if (left.name > right.name) return 1;
    return 0;
  });
}

function escapeContractMarkers(value: string): string {
  return value
    .replaceAll(CONTRACT_START, "&lt;!-- pi-orchestrate:contract:start --&gt;")
    .replaceAll(CONTRACT_END, "&lt;!-- pi-orchestrate:contract:end --&gt;");
}

function formatCatalog(catalog: WorkerCatalog): string {
  const workers = sortedWorkers(catalog);
  if (workers.length === 0) return "- No trusted workers are available for this session.";

  return workers
    .map(
      (worker) =>
        `- \`${escapeContractMarkers(worker.name)}\` [${worker.source.kind}] (${worker.lifecycle}): ${escapeContractMarkers(worker.description)}`,
    )
    .join("\n");
}

interface ContractMarker {
  readonly start: number;
  readonly end: number;
  readonly kind: "start" | "end";
}

function contractMarkers(prompt: string): ContractMarker[] {
  const markers: ContractMarker[] = [];
  for (const [value, kind] of [
    [CONTRACT_START, "start"],
    [CONTRACT_END, "end"],
  ] as const) {
    let offset = 0;
    while (offset < prompt.length) {
      const start = prompt.indexOf(value, offset);
      if (start < 0) break;
      markers.push({ start, end: start + value.length, kind });
      offset = start + value.length;
    }
  }
  return markers.sort((left, right) => left.start - right.start);
}

function removeContractMarkers(prompt: string): {
  readonly prompt: string;
  readonly insertionOffset?: number;
} {
  const markers = contractMarkers(prompt);
  if (markers.length === 0) return { prompt };

  const removed: Array<{ start: number; end: number }> = [];
  const stack: ContractMarker[] = [];
  for (const marker of markers) {
    if (marker.kind === "start") {
      stack.push(marker);
      continue;
    }
    const start = stack.pop();
    if (start && stack.length === 0) removed.push({ start: start.start, end: marker.end });
  }

  for (const marker of markers) {
    if (!removed.some((range) => marker.start >= range.start && marker.end <= range.end)) {
      removed.push({ start: marker.start, end: marker.end });
    }
  }
  removed.sort((left, right) => left.start - right.start);

  const insertionPoint = markers[0]!.start;
  let insertionOffset = 0;
  let cursor = 0;
  let cleaned = "";
  for (const range of removed) {
    if (range.start < cursor) continue;
    const retained = prompt.slice(cursor, range.start);
    cleaned += retained;
    if (range.start <= insertionPoint) insertionOffset = cleaned.length;
    cursor = range.end;
  }
  cleaned += prompt.slice(cursor);
  return { prompt: cleaned, insertionOffset };
}

function buildContract(catalog: WorkerCatalog): string {
  return `${CONTRACT_START}
## Pi Orchestrate Contract

You are the parent orchestrator and own the task end to end.

### Delegation

- Keep trivial or tightly coupled work in the parent. Delegate work that can proceed independently or benefit from independent judgment.
- Choose worker scopes and counts from the task. Treat workers or counts named by the user as a floor unless the user sets an exact cap.
- Each \`orchestrate\` call creates a fresh worker session. Multiple calls may use the same worker definition and identical instructions when independent judgments are useful. Do not vary briefs merely to make them appear different. Interactive follow-up instead continues one worker ID with its existing context.
- Give each worker a self-contained brief with its objective, context, paths and scope, forbidden actions, success criteria, and expected output. Workers do not receive the parent conversation.

### Parallel dispatch

- Form the complete wave before emitting any tool call.
- For one worker, make one fully briefed \`orchestrate\` call.
- For N workers where N > 1, make exactly one \`multi_tool_use.parallel\` call. Its \`tool_uses\` must contain exactly N \`functions.orchestrate\` entries and no other tools.
- If \`multi_tool_use.parallel\` is not present, emit all N \`orchestrate\` calls as native siblings in one assistant response.
- Never dispatch a multi-worker wave as separate assistant responses. An admitted sole asynchronous \`orchestrate\` call ends the parent turn, so omitted workers cannot be added afterward.
- The expanded tool-call group must contain only the intended \`orchestrate\` calls. Mixing another tool into the group makes orchestration inline and blocking.

### Completion and lifecycle

- Calls are admitted independently; a rejected call does not stop its siblings.
- After dispatching, wait for automatic result delivery instead of polling \`worker_status\`. When results expose more independent work, dispatch another complete wave.
- Automatic delivery requires no keepalive activity. While awaiting it, do not call \`sleep\`, poll with any tool, inspect files or processes to infer worker progress, or issue no-op tool calls. Perform only genuinely independent work that would be useful even if no worker were active; otherwise end the turn.
- The parent reviews and synthesizes worker results, resolves conflicts, integrates changes, and runs the relevant verification.
- Prefer one-shot workers. Use interactive workers only when retained context is useful, and follow the ownership and status requirements in the lifecycle tool descriptions.

### Trusted worker catalog

${formatCatalog(catalog)}
${CONTRACT_END}`;
}

export function applyOrchestratorContract(
  systemPrompt: string,
  catalog: WorkerCatalog,
): string {
  const section = buildContract(catalog);
  const cleaned = removeContractMarkers(systemPrompt);
  if (cleaned.insertionOffset !== undefined) {
    return `${cleaned.prompt.slice(0, cleaned.insertionOffset)}${section}${cleaned.prompt.slice(cleaned.insertionOffset)}`;
  }

  const separator =
    systemPrompt.length === 0 || systemPrompt.endsWith("\n\n")
      ? ""
      : systemPrompt.endsWith("\n")
        ? "\n"
        : "\n\n";
  return `${systemPrompt}${separator}${section}`;
}