/**
 * K3 Cache Optimization Extension for pi-agent — Phase 1 MVP
 *
 * Scope (handover doc §8 Phase 1): ChannelDetector + non-K3 bypass,
 * CacheTelemetry (usage recompute + dual-channel dispatch + cost/quota),
 * /cache panel, PrefixGuard (frozen tool snapshot, additive diff guard,
 * deterministic payload view, 4K-token segment-hash prefix check),
 * drift alert + adjacent-turn divergence localization.
 *
 * Hard rules: zero fork, public extension API only; full bypass for non-K3;
 * history append-only — this extension NEVER rewrites messages, it only
 * observes, normalizes tool ordering, and alerts.
 *
 * Deployment note (verified Phase 0): pi loads extensions in order
 * project dir → global dir → explicit configured paths, and hooks run in
 * load order. Register this file via explicit path (settings/`-e`) so its
 * before_provider_request is the LAST gate in the chain.
 */

import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { completeSimple } from "@earendil-works/pi-ai/compat";
import { Type } from "@sinclair/typebox";
import { createHash } from "node:crypto";
import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { join } from "node:path";

// ════════════════════════════════════════════════════════════════════════════
// CONFIG — every threshold is configurable; overridden by ~/.pi/agent/k3-cache.json
// ════════════════════════════════════════════════════════════════════════════

const CONFIG = {
  /** Master switch (D3): false disables everything. */
  enabled: true,
  /** Per-turn hit-rate threshold below which divergence diagnostics run (D12). */
  hitRateAlertThreshold: 0.7,
  /** Segment size in chars for rolling prefix hashes (~4K tokens ≈ 16K chars). */
  segmentSizeChars: 16 * 1024,
  /** Idle gap treated as probable cache TTL expiry (community: TTL 5–30 min). */
  idleTtlMs: 5 * 60 * 1000,
  /** Noise floor: ignore hit-rate math below this many prompt tokens. */
  minPromptTokensForStats: 1024,
  /** D6 — session-locked reasoning effort applied when pi has none set. */
  defaultEffort: "high",
  /** D9 — warn when subscription context exceeds this fraction of the 256K window. */
  subscriptionWarnFraction: 0.8,
  /** D9 — on k3-256k overflow/threshold compaction, upgrade to k3 instead of compacting. */
  overflowUpgrade: true,
  /** D9 — after an overflow upgrade, auto-send a follow-up so the aborted turn resumes. */
  overflowAutoContinue: true,
  /** D9 — upgrade target (the only 1M-unlocking id on the subscription endpoint). */
  upgradeModelId: "k3",
  /**
   * D8 — defer pi's threshold compaction while context is below this fraction of the window.
   * Higher = fewer prefix forks on mid-length coding sessions (more stable cacheRead).
   */
  highWaterFraction: 0.85,
  /** D8 — summary chain merge triggers: max layers / max chain tokens (≈4 chars per tok). */
  summaryChainMaxLayers: 5,
  summaryChainMaxTokens: 8000,
  /** D8 — run the Δ-summary generation ourselves (off → pi native compaction). */
  summaryChainEnabled: true,
  /**
   * Compat — who owns compaction handling. "k3cache" (default) enables D8/D9
   * (defer, overflow upgrade, Δ-summary chain). Set "external" when another
   * compaction-managing extension is installed (e.g. pi-goal): pi's runner is
   * first-result-wins on session_before_compact, so exactly one owner must act.
   * Telemetry/attribution stays on either way (session_compact always fires).
   */
  compactionOwnership: "k3cache" as "k3cache" | "external",
  /** D10 — prewarm on resume after this idle gap; never periodic keep-alive. */
  prewarmEnabled: true,
  prewarmIdleMs: 5 * 60 * 1000,
  /**
   * D15/D17 — tool-result externalization (on-demand gate execution layer).
   * Default OFF so arm-A toolset matches bare pi (fair A/B cost). When enabled
   * (k3-cache.json), uses minChars + onDemandTools; gate logic is always tested.
   */
  externalizerEnabled: false,
  /** D17 gate: externalize iff size > threshold AND tool is low-reuse (list below).
   * Compat tip: add big-result tools from other extensions here (e.g. "fetch",
   * "web_search" from pi-web-access, "subagent" from pi-subagents). */
  externalizeMinChars: 8192,
  onDemandTools: ["bash"] as string[],
  /** §3 subscription endpoint constraints. */
  toolSchemaLimitBytes: 15 * 1024,
  requestBodyLimitBytes: 2 * 1024 * 1024,
  /** K3 model ids per provider (D3; verified against pi model catalog v0.82.1). */
  k3Models: {
    moonshotai: ["kimi-k3"],
    "moonshotai-cn": ["kimi-k3"],
    "kimi-coding": ["k3", "k3-256k", "kimi-for-coding", "kimi-for-coding-highspeed"],
  } as Record<string, string[]>,
  /**
   * API list price ($/M tokens) used to impute cost on the subscription
   * channel, where pi reports cost 0 (verified Phase 0). Doc §3.
   */
  price: { cacheRead: 0.3, input: 3.0, output: 15.0 },
  /** Tool-set guard (D7): allow additive-only changes; explicit allowlist of tools permitted to change/vanish. */
  toolChangeAllowlist: [] as string[],
};

try {
  const home = process.env.HOME || process.env.USERPROFILE || "";
  const p = join(home, ".pi", "agent", "k3-cache.json");
  if (home && existsSync(p)) Object.assign(CONFIG, JSON.parse(readFileSync(p, "utf8")));
} catch {
  /* config override is best-effort */
}

/** Summary LLM call — indirection so the destructive-matrix harness can stub it. */
const llmComplete: typeof completeSimple = (...args: Parameters<typeof completeSimple>) =>
  ((globalThis as any).__k3cacheCompleteSimple ?? completeSimple)(...args);

// ════════════════════════════════════════════════════════════════════════════
// Types & pure helpers
// ════════════════════════════════════════════════════════════════════════════

type Channel = "moonshotai" | "moonshotai-cn" | "kimi-coding" | null;

type MissReason =
  | "first_turn"
  | "prefix_drift"
  | "model_switch"
  | "idle_ttl"
  | "compaction"
  | "branch_nav"
  | "system_rebuild"
  | "tool_change"
  | "external_miss"
  | "none";

interface TurnRecord {
  turn: number;
  input: number;
  cacheRead: number;
  output: number;
  hitRate: number;
  cost: number;
  missReason: MissReason;
}

function sha(data: string): string {
  return createHash("sha256").update(data).digest("hex").slice(0, 16);
}

/** D7 allowlist match: exact names or simple globs ("goal_*", "*_draft"). */
export function toolChangeAllowed(name: string, allowlist: string[] = CONFIG.toolChangeAllowlist): boolean {
  return allowlist.some((p) =>
    p.includes("*")
      ? new RegExp(
          "^" + p.split("*").map((s) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")).join(".*") + "$",
        ).test(name)
      : p === name,
  );
}

/** Imputed USD cost at list price (also used by /cache panel and tests). */
export function imputeTokenCost(
  u: { input: number; cacheRead: number; output: number },
  price: { cacheRead: number; input: number; output: number } = CONFIG.price,
): number {
  return (u.cacheRead * price.cacheRead + u.input * price.input + u.output * price.output) / 1e6;
}

/**
 * D17 on-demand gate (pure): externalize large low-reuse tool tails so later
 * turns do not re-bill them as miss-input. Exported for matrix / unit coverage.
 */
export function d17ShouldExternalize(opts: {
  textLen: number;
  toolName: string;
  onDemandTools: string[];
  minChars: number;
  overWatermark: boolean;
}): boolean {
  const toolName = String(opts.toolName ?? "");
  const toolBase = toolName.includes(":") ? toolName.split(":").pop()! : toolName;
  const lowReuse =
    opts.onDemandTools.includes(toolName) || opts.onDemandTools.includes(toolBase);
  const big = opts.textLen > opts.minChars;
  return (big && lowReuse) || (opts.overWatermark && big);
}

/**
 * Deterministic JSON: sorted keys, stable recursion. Skips `cache_control`
 * markers — pi-ai moves them between turns on the anthropic-messages path and
 * the kimi endpoint ignores them, so they must not count as prefix bytes.
 */
function stableStringify(obj: unknown): string {
  if (obj === null || obj === undefined) return "null";
  if (typeof obj !== "object") return JSON.stringify(obj);
  if (Array.isArray(obj)) return "[" + obj.map(stableStringify).join(",") + "]";
  const rec = obj as Record<string, unknown>;
  return (
    "{" +
    Object.keys(rec)
      .sort()
      .filter((k) => k !== "cache_control")
      .map((k) => JSON.stringify(k) + ":" + stableStringify(rec[k]))
      .join(",") +
    "}"
  );
}

/**
 * Build the prefix-critical view of a provider payload as an ordered list of
 * chunks (so a divergence offset can be mapped back to a message index).
 * Dual-channel: OpenAI completions (messages[], tools[].function) vs
 * Anthropic messages (system[] + messages[], tools[].name).
 */
function payloadChunks(payload: any): { label: string; text: string }[] {
  const chunks: { label: string; text: string }[] = [];
  const isAnthropic = payload.system !== undefined && Array.isArray(payload.messages);
  // Note: model id is deliberately NOT part of the chunk stream — switches are
  // tracked via model_select (and the 256k→k3 upgrade legitimately keeps the cache).
  if (isAnthropic && payload.system !== undefined)
    chunks.push({ label: "system", text: stableStringify(payload.system) });
  if (Array.isArray(payload.tools)) {
    for (const t of payload.tools) {
      const name = t?.function?.name ?? t?.name ?? "?";
      chunks.push({ label: `tool:${name}`, text: stableStringify(t) });
    }
  }
  if (Array.isArray(payload.messages)) {
    payload.messages.forEach((m: any, i: number) => {
      chunks.push({ label: `msg[${i}]:${m?.role ?? "?"}`, text: stableStringify(m) });
    });
  }
  return chunks;
}

/** Segment a concatenated chunk stream into fixed-size hashes + chunk boundaries. */
function segmentize(chunks: { label: string; text: string }[]): {
  hashes: string[];
  boundaries: { label: string; endOffset: number }[];
  totalChars: number;
} {
  let full = "";
  const boundaries: { label: string; endOffset: number }[] = [];
  for (const c of chunks) {
    full += c.text;
    boundaries.push({ label: c.label, endOffset: full.length });
  }
  const hashes: string[] = [];
  for (let i = 0; i < full.length; i += CONFIG.segmentSizeChars)
    hashes.push(sha(full.slice(i, i + CONFIG.segmentSizeChars)));
  return { hashes, boundaries, totalChars: full.length };
}

/** Map a char offset to the chunk label containing it. */
function chunkAtOffset(boundaries: { label: string; endOffset: number }[], offset: number): string {
  for (const b of boundaries) if (offset < b.endOffset) return b.label;
  return boundaries.length ? boundaries[boundaries.length - 1].label : "(empty)";
}

/**
 * Append-only prefix check over the segment-hash chain (doc §5: compare
 * chains, O(#segments)). Every full segment of the previous request except
 * its last (possibly partial, legitimately extended) one must reappear
 * unchanged at the same position. Returns violating segment index or -1.
 */
function prefixViolation(prev: string[], curr: string[]): number {
  const mustMatch = Math.min(prev.length - 1, curr.length);
  for (let i = 0; i < mustMatch; i++) if (prev[i] !== curr[i]) return i;
  if (curr.length < prev.length - 1) return curr.length; // history shrank
  return -1;
}

/**
 * Chunk-level append-only check: all previously sent chunks (model/system/
 * tool/message units) must reappear byte-identical at the same index — an
 * honest turn only appends. Catches drift inside the last partial segment
 * that the coarse chain misses; also yields the exact diverging unit label.
 */
function chunkViolation(
  prev: { label: string; hash: string }[],
  curr: { label: string; hash: string }[],
): number {
  const n = Math.min(prev.length, curr.length);
  for (let i = 0; i < n; i++) if (prev[i].hash !== curr[i].hash) return i;
  if (curr.length < prev.length) return curr.length; // history shrank
  return -1;
}

// ════════════════════════════════════════════════════════════════════════════
// Extension entry point
// ════════════════════════════════════════════════════════════════════════════

export default function k3CacheExtension(pi: ExtensionAPI) {
  /**
   * Per-session state, keyed by session id: in-process multi-session hosts
   * (pi-subagents and friends) must not interleave PrefixGuard hash chains
   * across sessions (that would produce false drift alerts). Hosts without
   * a session id all map to "__default" — identical to the old behavior.
   */
  function freshState() {
    return {
      active: false,
      channel: null as Channel,
      modelId: null as string | null,
      turn: 0,
      // PrefixGuard
      frozenTools: null as Map<string, string> | null, // name → schema hash
      prevHashes: [] as string[],
      prevBoundaries: [] as { label: string; endOffset: number }[],
      prevChunks: [] as { label: string; hash: string }[],
      prevChars: 0,
      driftThisTurn: false,
      driftDetail: "",
      driftAlerts: 0,
      toolGuardTrips: 0,
      lastToolGuardSig: "",
      lastToolGuardSigCount: 0,
      // Telemetry
      records: [] as TurnRecord[],
      totals: { input: 0, cacheRead: 0, output: 0, cost: 0 },
      lastRequestTs: 0,
      /** Next legitimate miss cause (set by lifecycle events, consumed once). */
      pendingReason: "none" as MissReason,
      lastSystemPrompt: "",
      // EffortManager (D6)
      lockedEffort: "" as string,
      effortSelfChange: false,
      // ModeManager (D9, subscription only)
      lastContextTokens: 0,
      quotaWarned: false,
      upgraded: false,
      // CompactionGuard (D8)
      summaryChain: [] as string[],
      pendingChain: null as string[] | null,
      // Externalizer (D15/D17)
      externalizedCount: 0,
      schemaWarned: false,
      // Prewarmer (D10)
      prewarmedAt: 0,
    };
  }
  type K3State = ReturnType<typeof freshState>;
  const states = new Map<string, K3State>();
  const sessionKey = (ctx: any): string => String(ctx?.sessionManager?.getSessionId?.() ?? "__default");
  function stateFor(ctx: any): K3State {
    const k = sessionKey(ctx);
    let s = states.get(k);
    if (!s) {
      s = freshState();
      states.set(k, s);
    }
    return s;
  }
  /** Feature-detected notify: hosts without ctx.ui must degrade, not crash. */
  const uiNotify = (ctx: any, msg: string, level?: "info" | "warning" | "error") =>
    ctx?.ui?.notify?.(msg, level);

  function detect(providerId: string | undefined, modelId: string | undefined): Channel {
    if (!CONFIG.enabled || !providerId || !modelId) return null;
    const ids = CONFIG.k3Models[providerId];
    return ids && ids.includes(modelId) ? (providerId as Channel) : null;
  }

  function applyModel(state: K3State, ctx: ExtensionContext, why: string) {
    const m = ctx.model;
    const ch = detect(m?.provider, m?.id);
    const wasActive = state.active;
    state.channel = ch;
    state.modelId = m?.id ?? null;
    state.active = ch !== null;
    if (state.active && !wasActive)
      uiNotify(ctx, `k3-cache: active (${ch}/${m!.id}, ${why})`, "info");
  }

  const imputedCost = (u: { input: number; cacheRead: number; output: number }) =>
    imputeTokenCost(u, CONFIG.price);

  // ── ChannelDetector ────────────────────────────────────────────────────────
  pi.on("session_start", (event, ctx) => {
    const state = freshState();
    states.set(sessionKey(ctx), state);
    applyModel(state, ctx, event.reason);
    if (!state.active) return;
    // D16 — compaction tuning advisory (settings are user-owned; warn, don't write).
    try {
      const home = process.env.HOME || process.env.USERPROFILE || "";
      const s = JSON.parse(readFileSync(join(home, ".pi", "agent", "settings.json"), "utf8"));
      const c = s?.compaction ?? {};
      if ((c.keepRecentTokens ?? 20000) < 100000 || (c.reserveTokens ?? 16384) < 32768)
        uiNotify(ctx, 
          `k3-cache: D16 tuning — recommend settings.json compaction {keepRecentTokens:100000, reserveTokens:32768} (current ${c.keepRecentTokens ?? "default"}/${c.reserveTokens ?? "default"}): fewer compactions = fewer prefix forks; K3 thinking counts as output and overflows 16K reserves`,
          "warning",
        );
    } catch {
      /* advisory only */
    }
    // D6 — lock the session's reasoning effort (pi defaults to settings; only
    // fill in when unset). Lock means "warn on change", never hard-forbid.
    state.lockedEffort = String(ctx.thinkingLevel ?? "");
    if (!state.lockedEffort && typeof ctx.setThinkingLevel === "function") {
      state.effortSelfChange = true;
      try {
        ctx.setThinkingLevel(CONFIG.defaultEffort as any);
        state.lockedEffort = CONFIG.defaultEffort;
      } finally {
        state.effortSelfChange = false;
      }
    }
    // Resume/fork: server cache may have expired or the active token
    // sequence changed — first miss is expected, attribute it correctly (D12).
    if (event.reason === "resume") state.pendingReason = "idle_ttl";
    else if (event.reason === "fork") state.pendingReason = "branch_nav";

    // Restore the summary chain of this session file (append-only custom entries).
    try {
      for (const e of ctx.sessionManager.getEntries() as any[]) {
        if (e?.type === "custom" && e.customType === "k3cache-chain") state.summaryChain = e.data?.chain ?? state.summaryChain;
      }
    } catch {
      /* chain restore is best-effort */
    }

    // ── Prewarmer (D10): resume from disk + idle beyond TTL → one cheap request
    // through the same pi-ai serialization path re-warms the server prefix.
    if (CONFIG.prewarmEnabled && (event.reason === "resume" || event.reason === "startup")) {
      try {
        const entries = ctx.sessionManager.getEntries() as any[];
        const msgs = entries.filter((e) => e?.type === "message" && e.message).map((e) => e.message);
        const last = entries[entries.length - 1];
        const lastTs = last ? Date.parse(last.timestamp ?? "") || 0 : 0;
        if (msgs.length >= 2 && lastTs && Date.now() - lastTs > CONFIG.prewarmIdleMs) {
          const m = ctx.model!;
          void (async () => {
            try {
              const auth = await ctx.modelRegistry.getApiKeyAndHeaders(m);
              await llmComplete(m as any, { systemPrompt: ctx.getSystemPrompt(), messages: msgs as any }, {
                apiKey: (auth as any)?.apiKey,
                maxTokens: 1,
              } as any);
              state.prewarmedAt = Date.now();
              pi.appendEntry("k3cache-prewarm", { messages: msgs.length });
              uiNotify(ctx, "k3-cache: prewarmed resumed session (server prefix re-cached)", "info");
            } catch {
              /* prewarm is opportunistic; a failure just means the first real request pays the miss */
            }
          })();
        }
      } catch {
        /* prewarm detection best-effort */
      }
    }
  });

  pi.on("model_select", (event, ctx) => {
    const state = stateFor(ctx);
    const before = state.active;
    applyModel(state, ctx, "model_select");
    if (before && state.active && event.previousModel && event.previousModel.id !== event.model.id) {
      // Official exception (D11): k3-256k → k3 upgrade keeps the cache.
      const upgrade = state.channel === "kimi-coding" && event.previousModel.id === "k3-256k" && event.model.id === "k3";
      if (!upgrade) {
        state.pendingReason = "model_switch";
        uiNotify(ctx, "k3-cache: model switched — server prefix cache will rebuild (expected miss)", "warning");
      }
    }
  });

  pi.on("session_compact", (_event, ctx) => {
    const state = stateFor(ctx);
    if (state.active) {
      state.pendingReason = "compaction";
      // Compaction legitimately rewrites the request prefix; tool snapshot survives.
      state.prevHashes = [];
      state.prevBoundaries = [];
      state.prevChunks = [];
      // Commit the summary chain layer (append-only custom entry, D8).
      if (state.pendingChain) {
        state.summaryChain = state.pendingChain;
        state.pendingChain = null;
        pi.appendEntry("k3cache-chain", { chain: state.summaryChain, layers: state.summaryChain.length });
      }
    }
  });

  // ── EffortManager (D6): session lock + warned switches ────────────────────
  // Matrix ⑤: reasoning level is conservatively treated as part of the cache
  // key (official side of the §6.2 disagreement). Warn, attribute, let through.
  pi.on("thinking_level_select", (event, ctx) => {
    const state = stateFor(ctx);
    if (!state.active || event.level === event.previousLevel) return;
    if (state.effortSelfChange) return; // our own /effort or session-start init
    state.pendingReason = "model_switch";
    state.lockedEffort = String(event.level);
    uiNotify(ctx, 
      `k3-cache: reasoning effort ${event.previousLevel} → ${event.level} — official docs say the server cache rebuilds (expected miss). Locked level updated; prefer one effort per session (D6)`,
      "warning",
    );
  });

  pi.registerCommand("effort", {
    description: "K3: switch session-locked reasoning effort (warns: cache rebuilds)",
    getArgumentCompletions: (prefix: string) =>
      ["minimal", "low", "medium", "high", "xhigh"]
        .filter((l) => l.startsWith(prefix))
        .map((l) => ({ value: l, label: l })),
    handler: async (args, ctx) => {
      const state = stateFor(ctx);
      if (!state.active) {
        uiNotify(ctx, "k3-cache: bypassed (not a K3 model)", "info");
        return;
      }
      const level = args.trim();
      if (!level) {
        uiNotify(ctx, `k3-cache: session effort locked at "${state.lockedEffort || ctx.thinkingLevel}" (D6). Usage: /effort <level>`, "info");
        return;
      }
      if (level === state.lockedEffort) {
        uiNotify(ctx, `k3-cache: effort already "${level}"`, "info");
        return;
      }
      const ok = ctx.ui?.confirm
        ? await ctx.ui.confirm(
            "K3 cache will rebuild",
            `Switching reasoning effort (${state.lockedEffort} → ${level}) invalidates the server prefix cache per official docs — the next request re-bills the full prompt at 10× the hit price. Switch anyway?`,
          )
        : true;
      if (!ok) return;
      state.effortSelfChange = true;
      try {
        if (typeof ctx.setThinkingLevel === "function") ctx.setThinkingLevel(level as any);
      } finally {
        state.effortSelfChange = false;
      }
      const from = state.lockedEffort;
      state.lockedEffort = level;
      state.pendingReason = "model_switch";
      pi.appendEntry("k3cache-effort-change", { from, to: level, turn: state.turn });
      uiNotify(ctx, `k3-cache: effort → ${level}; expect one full-prompt miss, then hits resume`, "warning");
    },
  });

  /**
   * K3-specific system-prompt constraints, appended to pi's native system
   * prompt (base is preserved — tools / environment / safety rules live there
   * and K3 is sensitive to framework context). Byte-identical every turn →
   * prefix-safe. Scope-first by design: K3's documented failure mode is
   * excessive proactiveness, and Moonshot explicitly advises "impose more
   * explicit behavioral constraints on K3 in the system prompt". So we BOUND
   * the task instead of over-instructing a frontier-class model. The XML tag
   * doubles as the stable idempotency anchor (never re-append within a turn).
   */
  const K3_SYSTEM_CONSTRAINTS =
    "\n\n<k3-cache-constraints>\n" +
    "You are a frontier-class coding model; scope discipline matters more than instruction volume. Stay inside the task and resist over-acting on ambiguity.\n" +
    "- Scope: change only what the issue requires. Do not refactor, rename, restyle, or edit unrelated code; do not change a working value's type or an API's shape for consistency alone. When intent is ambiguous, take the narrowest reading rather than expanding it.\n" +
    "- Completeness: the same behavior often lives in several places (a base class and its subclasses, a backend and its form layer, parallel handlers). Before finishing, search for sibling implementations of what you changed and fix every affected one — covering only one of several affected files is not smaller, it is incomplete.\n" +
    "- Verification: confirm by reading code plus at most one targeted snippet; do not run the full suite. If asked to re-check work that already passed, change it only when you can name a concrete failing input, otherwise confirm it and stop.\n" +
    "- Economy: reuse what is already in the transcript. Do not re-read files or re-dump large output already shown, and do not search full-repository git history — the issue text and current code are sufficient.\n" +
    "</k3-cache-constraints>";

  pi.on("before_agent_start", (event, ctx) => {
    const state = stateFor(ctx);
    if (!state.active) return;
    const base = String((event as any).systemPrompt ?? "");
    if (!base || base.includes("<k3-cache-constraints>")) return;
    return { systemPrompt: base + K3_SYSTEM_CONSTRAINTS };
  });

  // ── ModeManager (D9, subscription channel) ─────────────────────────────
  async function upgradeTo1M(state: K3State, ctx: ExtensionContext, why: string): Promise<boolean> {
    const target = (ctx.modelRegistry as any)?.find?.("kimi-coding", CONFIG.upgradeModelId);
    if (!target || typeof ctx.setModel !== "function") {
      uiNotify(ctx, `k3-cache: upgrade failed — model kimi-coding/${CONFIG.upgradeModelId} not in registry`, "error");
      return false;
    }
    // Model switch goes through pi's model mechanism (§5: never swap ids in the
    // payload — that would corrupt pi's cost telemetry).
    const ok = await ctx.setModel(target);
    if (!ok) {
      uiNotify(ctx, "k3-cache: upgrade failed — no auth for target model", "error");
      return false;
    }
    state.upgraded = true;
    uiNotify(ctx, 
      `k3-cache: upgraded k3-256k → ${CONFIG.upgradeModelId} (1M) — ${why}. Server cache is KEPT (official exception); quota burn is ≈2× from here on`,
      "warning",
    );
    pi.appendEntry("k3cache-mode-upgrade", { why, turn: state.turn, contextTokens: state.lastContextTokens });
    return true;
  }

  pi.registerCommand("k3upgrade", {
    description: "K3 subscription: upgrade k3-256k → k3 (1M window, cache kept, ≈2× quota)",
    handler: async (_args, ctx) => {
      const state = stateFor(ctx);
      if (!state.active || state.channel !== "kimi-coding") {
        uiNotify(ctx, "k3-cache: /k3upgrade only applies on the kimi-coding channel", "info");
        return;
      }
      if (state.modelId !== "k3-256k") {
        uiNotify(ctx, `k3-cache: current model is ${state.modelId}, nothing to upgrade`, "info");
        return;
      }
      await upgradeTo1M(state, ctx, "manual /k3upgrade");
    },
  });

  // D9 overflow handling: on k3-256k, replace pi's native compact+retry with a
  // cache-preserving upgrade. Manual /compact is always respected.
  // D8 (all channels): defer below high water; otherwise summary chain — old
  // summaries are never rewritten, each compaction only APPENDS a Δ layer.
  pi.on("session_before_compact", async (event, ctx) => {
    const state = stateFor(ctx);
    if (!state.active) return;
    // Compat: pi's runner is first-result-wins here — when another extension
    // owns compaction (pi-goal etc.), stand down entirely (D8/D9 off).
    if (CONFIG.compactionOwnership === "external") return;

    // ── D9: subscription 256K → upgrade instead of compacting ──
    if (CONFIG.overflowUpgrade && state.channel === "kimi-coding" && state.modelId === "k3-256k" && event.reason !== "manual") {
      if (await upgradeTo1M(state, ctx, `context ${event.reason} on 256K window`)) {
        if (event.reason === "overflow" && event.willRetry && CONFIG.overflowAutoContinue) {
          // pi cancels its retry when we cancel the compaction; resume by appending
          // a follow-up user message at the tail (append-only → cache-safe).
          pi.sendUserMessage("Context window was upgraded to 1M (k3). Continue the interrupted task.", { deliverAs: "followUp" });
        }
        return { cancel: true };
      }
      // upgrade unavailable → fall through to summary-chain compaction
    }

    // ── D8: defer threshold compaction below the high-water mark ──
    const usage = typeof ctx.getContextUsage === "function" ? ctx.getContextUsage() : undefined;
    if (
      event.reason === "threshold" &&
      usage?.tokens != null &&
      usage.tokens < usage.contextWindow * CONFIG.highWaterFraction
    ) {
      uiNotify(ctx, 
        `k3-cache: deferred compaction (context ${Math.round((usage.percent ?? 0))}% < high water ${Math.round(CONFIG.highWaterFraction * 100)}%) — every compaction is a prefix fork`,
        "info",
      );
      return { cancel: true };
    }

    // ── D8: summary chain — Δₙ summarizes only the segment since the last
    // compaction point; prior layers are reproduced verbatim (append-only).
    if (!CONFIG.summaryChainEnabled) return; // pi native compaction
    try {
      const prep = event.preparation as any;
      const mustMerge =
        state.summaryChain.length >= CONFIG.summaryChainMaxLayers ||
        state.summaryChain.join("").length / 4 > CONFIG.summaryChainMaxTokens;
      const m = ctx.model!;
      const auth = await ctx.modelRegistry.getApiKeyAndHeaders(m as any);
      const transcript = (prep.messagesToSummarize as any[])
        .map((msg) => `[${msg.role}] ${typeof msg.content === "string" ? msg.content : JSON.stringify(msg.content)}`)
        .join("\n")
        .slice(0, 400_000);
      // Fixed prompt template (D8); summarization runs as an isolated request —
      // pi-ai sends no session affinity for these calls, matching pi's own
      // cacheRetention:"none" isolation for compaction.
      const deltaPrompt = mustMerge
        ? `Merge the following summary layers and new conversation segment into ONE concise summary (max 600 words). Preserve: task goal, key decisions, file paths touched, current state, next steps.\n\nEXISTING LAYERS:\n${state.summaryChain.join("\n\n")}\n\nNEW SEGMENT:\n${transcript}`
        : `Summarize ONLY the following conversation segment (max 300 words). Preserve: what was done, key decisions, file paths touched, current state, next steps.\n\nSEGMENT:\n${transcript}`;
      const res = await llmComplete(
        m as any,
        { systemPrompt: "You summarize coding-agent sessions. Output only the summary text.", messages: [{ role: "user", content: deltaPrompt, timestamp: Date.now() }] as any },
        { apiKey: (auth as any)?.apiKey, signal: event.signal } as any,
      );
      const delta = (res as any).content?.filter((c: any) => c.type === "text").map((c: any) => c.text).join("\n").trim();
      if (!delta) return; // fall back to native compaction
      const chain = mustMerge ? [delta] : [...state.summaryChain, delta];
      state.pendingChain = chain;
      const summary = chain.length === 1 ? chain[0] : chain.map((s, i) => `## Summary layer ${i + 1}\n${s}`).join("\n\n");
      return {
        compaction: {
          summary,
          firstKeptEntryId: prep.firstKeptEntryId,
          tokensBefore: prep.tokensBefore,
          usage: (res as any).usage,
          details: { k3SummaryChain: { layers: chain.length, merged: mustMerge } },
        },
      };
    } catch {
      // Any failure → let pi's native compaction handle it (never block).
      return;
    }
  });

  // ── PrefixGuard: byte-level gate (should be LAST in the hook chain) ───────
  pi.on("before_provider_request", (event, ctx) => {
    const state = stateFor(ctx);
    if (!state.active) return;
    const payload = event.payload as any;
    if (!payload || typeof payload !== "object") return;

    // ── 3.4 subscription endpoint constraints (§3) ──
    if (state.channel === "kimi-coding") {
      if (Array.isArray(payload.tools) && !state.schemaWarned) {
        for (const t of payload.tools) {
          const size = JSON.stringify(t).length;
          if (size > CONFIG.toolSchemaLimitBytes) {
            state.schemaWarned = true;
            uiNotify(ctx, `k3-cache: tool "${t?.name ?? "?"}" schema ${size}B exceeds the 15KB endpoint limit — requests may 400`, "warning");
          }
        }
      }
      // Foreign signed thinking blocks (from a non-kimi provider earlier in a
      // mixed session) break the /coding endpoint — strip them (§3). Kimi's own
      // blocks are never touched (reasoning replay is mandatory).
      try {
        const assistantProviders = (ctx.sessionManager.getEntries() as any[])
          .filter((e) => e?.type === "message" && e.message?.role === "assistant")
          .map((e) => e.message.provider);
        if (assistantProviders.some((p) => p && p !== "kimi-coding") && Array.isArray(payload.messages)) {
          let ai = 0;
          for (const msg of payload.messages) {
            if (msg?.role !== "assistant") continue;
            const provider = assistantProviders[ai++];
            if (provider && provider !== "kimi-coding" && Array.isArray(msg.content)) {
              const before = msg.content.length;
              msg.content = msg.content.filter((b: any) => !(b?.type === "thinking" && b.signature));
              if (msg.content.length < before)
                uiNotify(ctx, "k3-cache: stripped foreign-signed thinking block(s) from mixed-provider history (endpoint compat)", "warning");
            }
          }
        }
      } catch {
        /* provenance mapping is best-effort */
      }
      const bodySize = JSON.stringify(payload).length;
      if (bodySize > CONFIG.requestBodyLimitBytes)
        uiNotify(ctx, `k3-cache: request body ${(bodySize / 1048576).toFixed(2)}MB exceeds the 2MB endpoint limit — externalize or compact`, "warning");
    }

    const now = Date.now();
    if (state.lastRequestTs && now - state.lastRequestTs > CONFIG.idleTtlMs && state.pendingReason === "none")
      state.pendingReason = "idle_ttl";
    state.lastRequestTs = now;

    // D7 — deterministic tool ordering (in-place, additive-safe: same set, stable order).
    if (Array.isArray(payload.tools) && payload.tools.length > 1) {
      payload.tools.sort((a: any, b: any) =>
        String(a?.function?.name ?? a?.name ?? "").localeCompare(String(b?.function?.name ?? b?.name ?? "")),
      );
    }

    // D7 — frozen snapshot + additive-only diff.
    const toolMap = new Map<string, string>();
    if (Array.isArray(payload.tools))
      for (const t of payload.tools) toolMap.set(String(t?.function?.name ?? t?.name ?? "?"), sha(stableStringify(t)));
    if (!state.frozenTools) {
      state.frozenTools = toolMap;
    } else {
      const removed: string[] = [];
      const mutated: string[] = [];
      for (const [name, hash] of state.frozenTools) {
        if (toolChangeAllowed(name)) continue;
        if (!toolMap.has(name)) removed.push(name);
        else if (toolMap.get(name) !== hash) mutated.push(name);
      }
      if (removed.length || mutated.length) {
        state.toolGuardTrips++;
        // Attribute the resulting cache miss precisely (extensions like
        // pi-goal legitimately toggle their tools per turn — that rebuilds
        // the server prefix but is NOT history drift).
        if (state.pendingReason === "none") state.pendingReason = "tool_change";
        // Alert once per distinct change-set — a per-turn toggler would
        // otherwise spam an identical error every request.
        const sig = removed.join(",") + "|" + mutated.join(",");
        if (sig !== state.lastToolGuardSig) {
          state.lastToolGuardSig = sig;
          state.lastToolGuardSigCount = 1;
          uiNotify(
            ctx,
            `k3-cache: NON-ADDITIVE tool change breaks the prefix cache` +
              (removed.length ? ` | removed: ${removed.join(", ")}` : "") +
              (mutated.length ? ` | schema changed: ${mutated.join(", ")}` : "") +
              ` — additive-only is the rule; allowlist expected togglers via toolChangeAllowlist (globs ok, e.g. "goal_*") or start /new (D7)`,
            "error",
          );
        } else {
          state.lastToolGuardSigCount++;
          // Recurring toggler (goal gating, MCP metadata refresh …): upgrade
          // the detection to a concrete remedy, exactly once at the 3rd hit.
          if (state.lastToolGuardSigCount === 3)
            uiNotify(
              ctx,
              `k3-cache: the same tool change-set recurred ${state.lastToolGuardSigCount}× — an extension toggles tools by design. If expected, silence alerts (attribution stays) via ~/.pi/agent/k3-cache.json: {"toolChangeAllowlist":${JSON.stringify([...removed, ...mutated])}}`,
              "warning",
            );
        }
        // Snapshot follows reality so one bad change doesn't alert forever.
        state.frozenTools = toolMap;
      } else {
        let added = false;
        for (const [name, hash] of toolMap)
          if (!state.frozenTools.has(name)) {
            state.frozenTools.set(name, hash);
            added = true;
          }
        // Mid-session additive registration is legal (no alert) but still
        // rewrites the tools region of the prefix (tools precede messages) —
        // attribute the miss so it is not misreported as history drift.
        if (added && state.pendingReason === "none") state.pendingReason = "tool_change";
      }
    }

    // System prompt rebuild (tools/skills/AGENTS.md changed) is a legitimate,
    // attributable prefix change — not drift (§6.8).
    const sys = typeof ctx.getSystemPrompt === "function" ? ctx.getSystemPrompt() : "";
    if (state.lastSystemPrompt && sys !== state.lastSystemPrompt && state.pendingReason === "none")
      state.pendingReason = "system_rebuild";
    state.lastSystemPrompt = sys;

    // Segment-hash prefix verification (1.4/1.5): coarse chain (O(#segments))
    // plus chunk-level append-only check for precise localization and for
    // drift hiding inside the last partial segment.
    const chunks = payloadChunks(payload);
    const chunkHashes = chunks.map((c) => ({ label: c.label, hash: sha(c.text) }));
    const { hashes, boundaries, totalChars } = segmentize(chunks);
    state.driftThisTurn = false;
    state.driftDetail = "";
    if ((state.prevHashes.length || state.prevChunks.length) && state.pendingReason === "none") {
      const segBad = prefixViolation(state.prevHashes, hashes);
      const chkBad = chunkViolation(state.prevChunks, chunkHashes);
      if (segBad >= 0 || chkBad >= 0) {
        const offset = Math.max(segBad, 0) * CONFIG.segmentSizeChars;
        const where =
          chkBad >= 0
            ? (chunkHashes[chkBad] ?? state.prevChunks[chkBad]).label
            : chunkAtOffset(state.prevBoundaries, Math.min(offset, Math.max(state.prevChars - 1, 0)));
        state.driftThisTurn = true;
        state.driftAlerts++;
        state.driftDetail = `${segBad >= 0 ? `segment #${segBad} (~${Math.round(offset / 4)} tok) ` : ""}at ${where}`;
        uiNotify(ctx, 
          `k3-cache: PREFIX DRIFT — earlier bytes changed ${state.driftDetail}; ` +
            `history must be append-only (volatile content only at the tail)`,
          "error",
        );
      }
    }
    state.prevHashes = hashes;
    state.prevBoundaries = boundaries;
    state.prevChunks = chunkHashes;
    state.prevChars = totalChars;
    // Never replace the payload object; in-place normalization only.
    return undefined;
  });

  // ── CacheTelemetry ─────────────────────────────────────────────────────────
  pi.on("message_end", (event, ctx) => {
    const state = stateFor(ctx);
    if (!state.active) return;
    const msg = event.message as any;
    if (msg?.role !== "assistant" || !msg.usage) return;
    const u = msg.usage;
    const input = u.input ?? 0;
    const cacheRead = u.cacheRead ?? 0;
    const output = u.output ?? 0;
    const denom = input + cacheRead;
    state.turn++;
    const hitRate = denom > 0 ? cacheRead / denom : 0;
    // Subscription channel reports cost 0 (Phase 0 finding) → impute at API list price.
    const cost = u.cost?.total > 0 ? u.cost.total : imputedCost({ input, cacheRead, output });

    let reason: MissReason = "none";
    const lowHit = denom >= CONFIG.minPromptTokensForStats && hitRate < CONFIG.hitRateAlertThreshold;
    if (state.turn === 1 && state.pendingReason === "none") reason = lowHit ? "first_turn" : "none";
    else if (state.driftThisTurn) reason = "prefix_drift";
    else if (state.pendingReason !== "none") reason = lowHit ? state.pendingReason : "none";
    else if (lowHit) reason = "external_miss"; // identical prefix, still missed → provider-side (§7-⑧)
    state.pendingReason = "none";

    const rec: TurnRecord = { turn: state.turn, input, cacheRead, output, hitRate, cost, missReason: reason };
    state.records.push(rec);
    state.totals.input += input;
    state.totals.cacheRead += cacheRead;
    state.totals.output += output;
    state.totals.cost += cost;
    pi.appendEntry("k3cache-usage", { ...rec, channel: state.channel, model: state.modelId, drift: state.driftDetail || undefined });

    // ModeManager watermark (D9): context ≈ prompt + this output, vs 256K window.
    state.lastContextTokens = input + cacheRead + output;
    if (state.channel === "kimi-coding" && state.modelId === "k3-256k" && !state.quotaWarned) {
      const window = (ctx.model as any)?.contextWindow ?? 262144;
      if (state.lastContextTokens > window * CONFIG.subscriptionWarnFraction) {
        state.quotaWarned = true;
        uiNotify(ctx, 
          `k3-cache: context ${state.lastContextTokens} tok > ${Math.round(CONFIG.subscriptionWarnFraction * 100)}% of the 256K window — approaching the limit. /k3upgrade switches to k3 (1M): cache kept, quota ≈2×. On overflow the upgrade happens automatically`,
          "warning",
        );
      }
    }

    if (lowHit && reason !== "none" && reason !== "first_turn") {
      const label: Record<MissReason, string> = {
        first_turn: "first request of session",
        prefix_drift: `prefix drift (${state.driftDetail})`,
        model_switch: "model switch rebuilt the cache",
        idle_ttl: "idle beyond cache TTL",
        compaction: "compaction changed the prefix (expected once)",
        branch_nav: "branch navigation changed the active token sequence",
        system_rebuild: "system prompt rebuilt (tools/skills/AGENTS.md changed)",
        tool_change: "tool set changed non-additively (extension toggled tools) — prefix cache rebuilt",
        external_miss: "external miss — bytes identical, provider-side routing/eviction",
        none: "",
      };
      uiNotify(ctx, 
        `k3-cache: hit rate ${(hitRate * 100).toFixed(0)}% < ${(CONFIG.hitRateAlertThreshold * 100).toFixed(0)}% — ${label[reason]}`,
        "warning",
      );
    }
  });

  // ── OnDemandGate + tool-result size control (D15/D17) ────────────────────
  // Always-on soft path: no new tools (fair A/B toolset). Large low-reuse bash
  // dumps are frozen as head+tail at write time so later turns do not re-bill
  // mega tails as miss-input. Optional full externalizer adds read_result when
  // CONFIG.externalizerEnabled (stable registration at load only).
  if (CONFIG.externalizerEnabled) {
    const storeDir = () => {
      const d = join(process.cwd(), "result-store");
      mkdirSync(d, { recursive: true });
      return d;
    };

    pi.registerTool({
      name: "read_result",
      label: "Read externalized result",
      description:
        "Retrieve an externalized tool result by id (optionally a line range). Content returns at the tail of the conversation — cache-safe.",
      parameters: Type.Object({
        id: Type.String({ description: "Result id from the placeholder" }),
        startLine: Type.Optional(Type.Number()),
        endLine: Type.Optional(Type.Number()),
      }),
      async execute(_id, params) {
        const file = join(storeDir(), `${String(params.id).replace(/[^a-f0-9]/gi, "")}.txt`);
        if (!existsSync(file))
          return {
            content: [{ type: "text", text: `No externalized result ${params.id}` }],
            isError: true,
          } as any;
        let text = readFileSync(file, "utf8");
        if (params.startLine || params.endLine) {
          const lines = text.split("\n");
          text = lines
            .slice((params.startLine ?? 1) - 1, params.endLine ?? lines.length)
            .join("\n");
        }
        return { content: [{ type: "text", text }] } as any;
      },
    } as any);
  }

  pi.on("tool_result" as any, (event: any, ctx: ExtensionContext) => {
    const state = stateFor(ctx);
    if (!state.active || event.isError) return;
    const text = (event.content ?? [])
      .filter((c: any) => c.type === "text")
      .map((c: any) => c.text)
      .join("\n");
    const toolName = String(event.toolName ?? "");
    const window = (ctx.model as any)?.contextWindow ?? 262144;
    const overWatermark =
      state.channel === "kimi-coding" &&
      state.modelId === "k3-256k" &&
      state.lastContextTokens > window * CONFIG.subscriptionWarnFraction;
    if (
      !d17ShouldExternalize({
        textLen: text.length,
        toolName,
        onDemandTools: CONFIG.onDemandTools,
        minChars: CONFIG.externalizeMinChars,
        overWatermark,
      })
    )
      return;

    try {
      const lines = text.split("\n");
      const id = sha(text);

      if (CONFIG.externalizerEnabled) {
        const storeDir = () => {
          const d = join(process.cwd(), "result-store");
          mkdirSync(d, { recursive: true });
          return d;
        };
        writeFileSync(join(storeDir(), `${id}.txt`), text, "utf8");
        const head = lines.slice(0, 8).join("\n").slice(0, 600);
        const tail =
          lines.length > 12 ? "\n...\n" + lines.slice(-4).join("\n").slice(0, 300) : "";
        const placeholder =
          `[k3-cache externalized result]\n` +
          `id: ${id}\n` +
          `tool: ${toolName}\n` +
          `lines: ${lines.length}, chars: ${text.length}\n` +
          `preview:\n${head}${tail}\n\n` +
          `Use the read_result tool with id "${id}" only if you need the full body. ` +
          `Prefer acting on the preview when sufficient (avoids re-billing large tails).`;
        state.externalizedCount++;
        pi.appendEntry("k3cache-externalized", { id, tool: toolName, chars: text.length, mode: "store" });
        return { content: [{ type: "text", text: placeholder }] };
      }

      // Soft path (default): freeze head+tail only — no new tools, fair vs bare pi.
      const headN = 40;
      const tailN = 20;
      const head = lines.slice(0, headN).join("\n");
      const tail = lines.length > headN + tailN ? lines.slice(-tailN).join("\n") : "";
      const truncated =
        `[k3-cache truncated tool result | tool=${toolName} | lines=${lines.length} chars=${text.length} | id=${id}]\n` +
        head +
        (tail
          ? `\n\n...[${lines.length - headN - tailN} lines omitted to protect prefix-cache; re-run a narrower command if you need the middle]...\n\n` +
            tail
          : "") +
        `\n`;
      state.externalizedCount++;
      pi.appendEntry("k3cache-externalized", {
        id,
        tool: toolName,
        chars: text.length,
        mode: "truncate",
        kept: truncated.length,
      });
      return { content: [{ type: "text", text: truncated }] };
    } catch {
      return;
    }
  });

  // ── /cache panel ───────────────────────────────────────────────────────────
  pi.registerCommand("cache", {
    description: "K3 cache telemetry: hit rate, cost, drift alerts",
    handler: async (_args, ctx) => {
      const state = stateFor(ctx);
      if (!state.active) {
        uiNotify(ctx, "k3-cache: bypassed (current model is not Kimi K3)", "info");
        return;
      }
      const t = state.totals;
      const denom = t.input + t.cacheRead;
      const agg = denom ? (t.cacheRead / denom) * 100 : 0;
      const noCache = ((t.input + t.cacheRead) * CONFIG.price.input + t.output * CONFIG.price.output) / 1e6;
      const misses = state.records.filter((r) => r.missReason !== "none");
      const lines = [
        `K3 cache — ${state.channel}/${state.modelId}, ${state.turn} turns`,
        `hit rate: ${agg.toFixed(1)}% (read ${t.cacheRead} / miss-input ${t.input} / output ${t.output} tok)`,
        `cost${state.channel === "kimi-coding" ? " (imputed @API price)" : ""}: $${t.cost.toFixed(4)} | no-cache counterfactual: $${noCache.toFixed(4)}`,
        `drift alerts: ${state.driftAlerts} | tool-guard trips: ${state.toolGuardTrips}`,
        `effort lock: ${state.lockedEffort || "(none)"} | context: ${state.lastContextTokens} tok${state.upgraded ? " | mode: upgraded to 1M (≈2× quota)" : ""}`,
        `summary chain: ${state.summaryChain.length} layer(s) | prewarm: ${state.prewarmedAt ? "done" : "n/a"} | externalized: ${state.externalizedCount}`,
        misses.length
          ? `attributed misses: ${misses.map((r) => `#${r.turn} ${r.missReason}(${(r.hitRate * 100).toFixed(0)}%)`).join(", ")}`
          : `attributed misses: none`,
      ];
      // In-process multi-session hosts (subagents): aggregate across sessions.
      const tracked = [...states.values()].filter((s) => s.turn > 0);
      if (tracked.length > 1) {
        const all = tracked.reduce(
          (a, s) => ({ cost: a.cost + s.totals.cost, input: a.input + s.totals.input, cacheRead: a.cacheRead + s.totals.cacheRead }),
          { cost: 0, input: 0, cacheRead: 0 },
        );
        const aggHit = all.input + all.cacheRead ? (100 * all.cacheRead) / (all.input + all.cacheRead) : 0;
        lines.push(`all in-process sessions: ${tracked.length} tracked | total $${all.cost.toFixed(4)} | hit ${aggHit.toFixed(1)}%`);
      }
      uiNotify(ctx, lines.join("\n"), "info");
    },
  });
}