import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

/**
 * rename-pi — sanitize "pi" -> "the assistant" for the claude-proxy. (ToS violatooooor)
 *
 * WHY TWO HOOKS:
 *   The claude-proxy classifies requests by system-prompt content; bare "pi"
 *   references get billed as extra usage. Renaming must therefore reach EVERY
 *   outbound request.
 *
 *   `before_agent_start` only fires when a fresh user message is sent, and the
 *   system prompt it returns is later CLOBBERED whenever the prompt is rebuilt
 *   mid-run (e.g. pi-goal's create_goal -> setActiveTools -> _rebuildSystemPrompt
 *   regenerates the raw prompt with "pi"/"pi-goal" and the "Pi documentation"
 *   section). Goal continuations then ship ~40 un-sanitized "pi" refs and get
 *   billed as extra usage. That was the bug.
 *
 *   `before_provider_request` fires at the single choke point for ALL provider
 *   calls (typed turns, goal continuations, subagents, summarization), with the
 *   fully-built request payload, and can replace it. Sanitizing `payload.system`
 *   here is leak-proof by construction — no lifecycle rebuild can bypass it.
 *
 * We keep both: before_agent_start preserves the renamed prompt in-session/for
 * display; before_provider_request guarantees the wire is always sanitized.
 * The rewrite is idempotent ("the assistant" contains no bare "pi"), so running
 * both never double-mangles.
 *
 * Safety: only touches `payload.system` (string or Anthropic text-block array),
 * matching the original systemPrompt-only behavior. Other providers / shapes
 * are left untouched. Handler is defensive and never throws (the runner also
 * catches, but we belt-and-suspenders it).
 */

// Word-boundary-aware: only bare "pi"/"Pi" as standalone words. "the assistant"
// for grammatically correct output. Skips possessives / compounds mid-word.
function sanitize(text: string): string {
  return text.replace(/\bpi\b/g, "the assistant").replace(/\bPi\b/g, "The assistant");
}

// Sanitize the outbound request payload's system prompt in place (Anthropic
// `system` is an array of {type:"text", text} blocks; some APIs use a string).
// Returns a count of blocks changed, for optional debug.
function sanitizePayloadSystem(payload: any): number {
  let changed = 0;
  if (!payload || typeof payload !== "object") return changed;
  const sys = payload.system;
  if (typeof sys === "string") {
    const next = sanitize(sys);
    if (next !== sys) {
      payload.system = next;
      changed++;
    }
  } else if (Array.isArray(sys)) {
    for (const block of sys) {
      if (block && typeof block.text === "string") {
        const next = sanitize(block.text);
        if (next !== block.text) {
          block.text = next;
          changed++;
        }
      }
    }
  }
  return changed;
}

export default function (pi: ExtensionAPI) {
  // (1) In-session / display parity: rename the prompt shown to the model when a
  // fresh user turn starts. (Leaky on its own — see header — hence hook 2.)
  pi.on("before_agent_start", async (event, _ctx) => {
    return { systemPrompt: sanitize(event.systemPrompt) };
  });

  // (2) Wire guarantee: rewrite the system prompt on the actual outbound payload
  // for every provider request, regardless of which code path built it.
  pi.on("before_provider_request", async (event, _ctx) => {
    const payload = (event as any).payload;
    const changed = sanitizePayloadSystem(payload);
    if (changed && process.env.PI_SANITIZE_DEBUG) {
      try {
        process.stderr.write(`[rename-pi] sanitized ${changed} system block(s) on provider request\n`);
      } catch {
        /* ignore */
      }
    }
    return payload;
  });
}