/**
 * index.ts — Pi extension entrypoint.
 * Phase 3: observe + fingerprint + llama-server slot warm-up on miss.
 * Phase 5: disk persistence — on miss, restore the slot prompt cache from
 *   a checkpoint file (named by fingerprint) if one exists, verify it
 *   empirically, else cold-warm and save a new checkpoint (plan §14).
 *
 * Provider gate: the extension acts only for models whose provider
 *   baseUrl (resolved from the agent's models.json) equals
 *   PI_PREFIX_CACHE_BASE_URL — it can only warm its own server. Any
 *   other model (remote providers, other local servers) is skipped
 *   (log line only), including the slot-poll diagnostic.
 *
 * Flow per request:
 *   0. Provider gate: skip if the model is not served by BASE_URL.
 *   1. Extract stable prefix (sections 1–5) from the `developer` message.
 *   2. Compute base fingerprint (model, pi/ext versions, tools, kwargs,
 *      stable prefix).
 *   3. If this process already warmed this fingerprint → do nothing
 *      (the slot retains the KV; the normal request reuses it).
 *   4. Otherwise (miss), with persistence enabled:
 *        a. Try restoring slot prompt cache from `pi-prefix-<fp16>.bin`.
 *        b. Verify with a warm-up request: identical-prompt reuse must be
 *           near-complete (cache_n ≥ 90% of prompt_n). Invalid → erase +
 *           cold re-warm + re-save.
 *        c. No file / restore failed → cold warm-up, then save checkpoint.
 *   5. Let the normal Pi request proceed.
 *
 * Configuration (env):
 *   PI_PREFIX_CACHE_BASE_URL  llama-server base URL, e.g.
 *                             http://localhost:8081/v1
 *                             (unset → warm-up disabled, observe-only)
 *   PI_PREFIX_CACHE_SLOT      slot id for the shared prefix (default 0)
 *   PI_PREFIX_CACHE_PERSIST   "1" → enable disk save/restore (default off)
 *   PI_PREFIX_CACHE_SLOT_DIR  server's --slot-save-path; enables stale
 *                             checkpoint cleanup (Phase 6)
 *   PI_PREFIX_CACHE_CLEANUP   "0" to disable cleanup (default on)
 *   PI_PREFIX_CACHE_TTL_DAYS  keep checkpoints used within N days (30)
 *   PI_PREFIX_CACHE_KEEP_RECENT  always keep N most recent (3)
 *   PI_PREFIX_CACHE_MAX_GB  max total checkpoint storage in GB (0 = off);
 *                           oldest kept files (never the current one)
 *                           are deleted until under the cap
 *
 * Fail-open: any error is logged and the normal request proceeds
 * untouched. The payload is never modified (option A, findings §4).
 */
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { createRequire } from "node:module";
import { appendFileSync, readFileSync, realpathSync, statSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { extractStablePrefix, findSystemMessage } from "./prefix";
import { computeFingerprint } from "./fingerprint";
import { eraseSlotFile, getSlotStats, probeSlot, restoreSlotFile, saveSlotFile, warmSlot } from "./warmup";
import { cleanupStale, recordUse } from "./cache-store";

const EXTENSION_VERSION = "1.2.3";

/** Integer env var with fallback to a default when unset/invalid. */
function numEnv(name: string, fallback: number): number {
  const n = Number.parseInt(process.env[name] ?? "", 10);
  return Number.isFinite(n) ? n : fallback;
}

const BASE_URL = process.env.PI_PREFIX_CACHE_BASE_URL ?? "";
const SLOT = numEnv("PI_PREFIX_CACHE_SLOT", 0);
const PERSIST = (process.env.PI_PREFIX_CACHE_PERSIST ?? "") === "1";
/** Server's `--slot-save-path`; when set, stale pi-prefix-*.bin files
 *  are pruned once per process (Phase 6). */
const SLOT_DIR = process.env.PI_PREFIX_CACHE_SLOT_DIR ?? "";
const CLEANUP = (process.env.PI_PREFIX_CACHE_CLEANUP ?? "1") === "1";
const TTL_DAYS = numEnv("PI_PREFIX_CACHE_TTL_DAYS", 30);
const KEEP_RECENT = numEnv("PI_PREFIX_CACHE_KEEP_RECENT", 3);
/** Max total checkpoint storage in GB; 0 = uncapped. Oldest kept files
 *  (never the current session's) are deleted until under the cap. */
const MAX_GB = numEnv("PI_PREFIX_CACHE_MAX_GB", 0);
/** Verify a restored checkpoint only if ≥ this fraction of the warm-up
 *  prompt was served from cache (identical-prompt reuse is unbounded, so
 *  a valid restore must score near 1.0). */
const VERIFY_THRESHOLD = 0.9;
/** Optional llama-server build/runtime identity, included in the
 *  fingerprint. The server exposes no JSON build endpoint, so this is
 *  env-provided (e.g. `build-1676-slckv2`); bump it when you deploy a
 *  new llama-server build. */
const SERVER_ID = process.env.PI_PREFIX_CACHE_SERVER_ID ?? "";

/** Map model id → provider baseUrl from the agent's models.json (the
 *  same file pi loads; $PI_CODING_AGENT_DIR or ~/.pi/agent). Cached per
 *  process. Empty when the file is missing/unreadable — in that case the
 *  provider gate skips everything, which is the safe direction (never
 *  warm a server we cannot attribute the request to). */
let modelBaseUrls: Map<string, string> | null = null;
function modelBaseUrlsMap(): Map<string, string> {
  if (modelBaseUrls) return modelBaseUrls;
  const map = new Map<string, string>();
  try {
    const agentDir =
      (process.env.PI_CODING_AGENT_DIR ?? "").replace(/^~(?=\/|$)/, homedir()) ||
      join(homedir(), ".pi", "agent");
    const raw = JSON.parse(readFileSync(join(agentDir, "models.json"), "utf8")) as {
      providers?: Record<string, { baseUrl?: string; models?: Array<{ id?: string }> }>;
    };
    for (const prov of Object.values(raw.providers ?? {})) {
      if (typeof prov?.baseUrl !== "string" || !prov.baseUrl) continue;
      for (const m of prov.models ?? []) {
        if (typeof m?.id === "string") map.set(m.id, prov.baseUrl);
      }
    }
  } catch (err) {
    log(`[prefix-cache] WARN: cannot read models.json (${String(err)}) — provider gate will skip all requests`);
  }
  modelBaseUrls = map;
  return map;
}

/** True when the request's model is served by the provider whose
 *  baseUrl equals PI_PREFIX_CACHE_BASE_URL (trailing-slash-insensitive).
 *  The extension can only warm ITS OWN server, so any other model
 *  (remote providers, other local servers) must be skipped. */
function isTargetModel(model: string): boolean {
  if (!model) return false;
  const baseUrl = modelBaseUrlsMap().get(model);
  if (baseUrl === undefined) return false;
  return baseUrl.replace(/\/+$/, "") === BASE_URL.replace(/\/+$/, "");
}

/** Physical model identity: path + size + mtime; path only when the
 *  file cannot be stat'ed (relative path, gone, unreadable). */
function modelIdentity(model: string): string {
  try {
    if (model.startsWith("/")) {
      const st = statSync(model);
      return `${model}|size=${st.size}|mtime=${st.mtimeMs}`;
    }
  } catch {
    /* fall through to path only */
  }
  return model;
}

function piVersion(): string {
  try {
    const req = createRequire(import.meta.url);
    return String(req("@earendil-works/pi-coding-agent/package.json").version);
  } catch {
    /* fall through */
  }
  try {
    let entry = process.argv[1] ?? "";
    try {
      entry = realpathSync(entry);
    } catch {
      /* keep raw path */
    }
    const marker = "node_modules/@earendil-works/pi-coding-agent/";
    const idx = entry.indexOf(marker);
    if (idx !== -1) {
      const root = entry.slice(0, idx + marker.length - 1);
      return String(JSON.parse(readFileSync(root + "/package.json", "utf8")).version);
    }
  } catch {
    /* fall through */
  }
  return "unknown";
}

/** Diagnostics log file (default ~/.cache/pi/llama-prefix/debug.log;
 *  PI_PREFIX_CACHE_LOG overrides, empty string = disable). pi >= 1.0 runs
 *  the TUI fullscreen by default, so raw console.log writes land on top of
 *  the input area — log to a file instead. console.log only when
 *  PI_PREFIX_CACHE_DEBUG=1 (e.g. non-interactive runs). */
const LOG_FILE =
  process.env.PI_PREFIX_CACHE_LOG !== undefined
    ? process.env.PI_PREFIX_CACHE_LOG
    : join(homedir(), ".cache", "pi", "llama-prefix", "debug.log");
const DEBUG_CONSOLE = (process.env.PI_PREFIX_CACHE_DEBUG ?? "") === "1";
const log = (line: string) => {
  if (DEBUG_CONSOLE) console.log(line);
  if (!LOG_FILE) return;
  try {
    appendFileSync(LOG_FILE, `${new Date().toISOString()} ${line}\n`);
  } catch {
    /* fail-open: logging must never break the request */
  }
};

export default async function (pi: ExtensionAPI) {
  const version = piVersion();
  /** Fingerprint warmed in this process; null = nothing warmed yet. */
  let warmedFingerprint: string | null = null;
  /** Phase 6: stale-cache cleanup runs at most once per process. */
  let cleaned = false;
  /** Slot's n_prompt_tokens at our last observation; the after-response
   *  hook waits for it to change, which means the normal request's
   *  prompt has been processed (slot counters reflect the last prompt). */
  let lastSeenPromptTokens = -1;
  /** The last before-request passed the provider gate (its traffic goes
   *  to BASE_URL); the after-response hook polls the slot only for such
   *  requests (the event carries no model, so state comes from before). */
  let lastWasTarget = false;
  /** One-shot TUI notification: the first target request of this process
   *  announces the cache outcome (WARM/RESTORED/COLD); later requests are
   *  silent (file log only). pi >= 1.0 runs fullscreen, so console.log
   *  would clobber the input area — ctx.ui.notify renders in-TUI. */
  let notified = false;

  pi.on("before_provider_request", async (event, ctx) => {
    const notify = (message: string, type?: "info" | "warning" | "error") => {
      if (notified) return;
      notified = true;
      try {
        ctx.ui.notify(message, type);
      } catch {
        /* non-interactive mode / no UI — file log already has it */
      }
    };
    try {
      const payload = event.payload as {
        model?: string;
        messages?: Array<{ role?: string; content?: unknown }>;
        tools?: unknown;
        chat_template_kwargs?: unknown;
      };

      // ---- Provider gate: act only for models served by BASE_URL ----
      // The event carries no provider; resolve model → provider baseUrl
      // from models.json. Remote models (and local servers other than
      // BASE_URL) must not trigger warm-up or slot polling.
      lastWasTarget = !BASE_URL || isTargetModel(payload.model ?? "");
      if (BASE_URL && !lastWasTarget) {
        log(`[prefix-cache] skip: model "${payload.model ?? ""}" is not served by ${BASE_URL} (provider base URL mismatch)`);
        return;
      }

      const sys = findSystemMessage(payload.messages ?? []);
      if (!sys) {
        log("[prefix-cache] skip: no developer/system message in payload");
        return;
      }

      const sp = extractStablePrefix(sys.content);
      if (!sp.ok) {
        log(`[prefix-cache] skip (fail-open): ${sp.reason}`);
        return;
      }

      const tools = payload.tools ?? [];
      const toolCount = Array.isArray(tools) ? tools.length : -1;

      const fp = computeFingerprint({
        model: payload.model ?? "",
        modelIdentity: modelIdentity(payload.model ?? ""),
        serverId: SERVER_ID,
        piVersion: version,
        extensionVersion: EXTENSION_VERSION,
        stablePrefix: sp.text,
        tools,
        chatTemplateKwargs: payload.chat_template_kwargs ?? {},
      });

      log(`[prefix-cache] model=${payload.model}`);
      log(`[prefix-cache] pi=${version} ext=${EXTENSION_VERSION} sys-role=${sys.role}`);
      log(`[prefix-cache] boundary=${sp.boundary}@${sp.boundaryIndex} stable-chars=${sp.text.length}`);
      log(`[prefix-cache] tools=${toolCount} kwargs=${JSON.stringify(payload.chat_template_kwargs ?? {})}`);
      log(`[prefix-cache] fingerprint=${fp}`);

      // ---- Phase 6: one-time stale-cache cleanup (fail-open) ----
      if (!cleaned && PERSIST && CLEANUP && SLOT_DIR) {
        cleaned = true;
        try {
          const c = cleanupStale(SLOT_DIR, [fp], {
            ttlDays: TTL_DAYS,
            keepRecent: KEEP_RECENT,
            maxBytes: MAX_GB > 0 ? MAX_GB * 1e9 : undefined,
          });
          if (c.skipped) {
            log(`[prefix-cache] cleanup SKIPPED (fail-safe, index untrusted): ${c.skipped}`);
          } else if (c.deleted.length > 0) {
            log(
              `[prefix-cache] cleanup: deleted ${c.deleted.length} stale checkpoint(s), ` +
                `freed ${(c.totalBytesDeleted / 1e9).toFixed(2)} GB`,
            );
          }
        } catch (err) {
          log(`[prefix-cache] cleanup failed (fail-open): ${String(err)}`);
        }
      }

      // ---- Phase 3/5: slot warm-up (+ optional disk save/restore) ----
      if (!BASE_URL) {
        log("[prefix-cache] warm-up disabled (PI_PREFIX_CACHE_BASE_URL unset)");
        return;
      }

      if (warmedFingerprint === fp) {
        // Liveness (review fix #1): a server restart clears the slot
        // while this process still believes it is warm. A fresh slot
        // has a zeroed prompt counter, so trust the in-memory state
        // only if the slot is alive. Unreachable server → fail open
        // (the normal request would fail anyway).
        let live = true;
        try {
          const s = await getSlotStats(BASE_URL, SLOT);
          live = !!s && s.n_prompt_tokens > 0;
        } catch {
          /* keep in-memory state */
        }
        if (live) {
          log(`[prefix-cache] cache=HIT (slot ${SLOT} already warmed for this fingerprint)`);
          notify("⚡ prefix-cache: WARM — slot already hot", "info");
          return;
        }
        log(`[prefix-cache] warmedFingerprint stale (slot ${SLOT} empty — server restarted?) → restore/re-warm`);
        warmedFingerprint = null;
      }

      const fname = `pi-prefix-${fp.slice(0, 16)}.bin`;
      // The Qwen3.8 template renders the tool schemas BEFORE the system
      // content, so the warm request must carry the same tools array or
      // the rendered prefixes diverge at token ~3 (findings: lcp = 3).
      const warm = () =>
        warmSlot(BASE_URL, payload.model ?? "", sys.role, sp.text, payload.chat_template_kwargs ?? {}, SLOT, undefined, tools);

      /** Phase 6 metadata fields for the index entry. */
      const meta = {
        fileName: fname,
        model: payload.model ?? "",
        baseUrl: BASE_URL,
        slot: SLOT,
        piVersion: version,
        extVersion: EXTENSION_VERSION,
      };

      let restored = false;

      if (PERSIST) {
        // (a) try restoring a checkpoint for this fingerprint
        const r = await restoreSlotFile(BASE_URL, SLOT, fname);
        if (r.ok && (r.n_restored ?? 0) > 0) {
          log(`[prefix-cache] RESTORE ${fname}: n_restored=${r.n_restored} (${r.timings?.restore_ms?.toFixed(1) ?? "?"} ms)`);
          // (b) verify empirically (plan §14): re-send the prefix; a valid
          // restore must serve it almost entirely from cache.
          // (b1) identical-prompt sanity gate: raw KV must be present.
          const v = await warm();
          const vn = v.timings?.prompt_n ?? 0;
          const vc = v.timings?.cache_n ?? 0;
          if (vn > 0 && vc >= VERIFY_THRESHOLD * vn) {
            // (b2) divergent probe (review fix #2): identical-prompt
            // reuse is a plain KV token match and does not prove the
            // restored chain can resume before the divergence point.
            // This probe diverges right after the stable prefix and
            // must be served from a context checkpoint (cache_n > 0).
            const p = await probeSlot(BASE_URL, payload.model ?? "", sys.role, sp.text, payload.chat_template_kwargs ?? {}, SLOT, tools);
            const pn = p.timings?.prompt_n ?? 0;
            const pc = p.timings?.cache_n ?? 0;
            if (pn > 0 && pc > 0) {
              log(`[prefix-cache] RESTORE verified: identical ${vc}/${vn}, divergent probe ${pc}/${pn}`);
            } else {
              // No checkpoint resume. A rebuild produces the same
              // chain, so this is almost always a server-flags issue
              // (e.g. prefix shorter than --checkpoint-min-step) —
              // keep the restore, warn loudly, proceed.
              log(
                `[prefix-cache] RESTORE probe FAILED (divergent cache_n=${pc}/${pn}): no checkpoint resume — ` +
                  `check --checkpoint-min-step/-ub; proceeding (full re-prefill expected)`,
              );
            }
            restored = true;
            notify(
              `⚡ prefix-cache: WARM — restored ${r.n_restored} tok from disk ` +
                `(${r.timings?.restore_ms?.toFixed(0) ?? "?"} ms)`,
              "info",
            );
            try {
              recordUse(fp, { ...meta, tokenCount: vn });
            } catch (err) {
              log(`[prefix-cache] metadata update failed (fail-open): ${String(err)}`);
            }
          } else {
            log(`[prefix-cache] RESTORE INVALID (verify cache_n=${vc}/${vn}) → erase + rebuild`);
            const e = await eraseSlotFile(BASE_URL, SLOT);
            if (e.ok) log(`[prefix-cache] slot erased (n_erased=${e.n_erased})`);
          }
        } else {
          log(`[prefix-cache] no checkpoint (restore HTTP ${r.status}${r.error ? `: ${r.error}` : ""}) → cold build`);
        }
      }

      if (!restored) {
        log(`[prefix-cache] cache=MISS → warming slot=${SLOT}${PERSIST ? " (cold)" : ""}`);
        const before = await getSlotStats(BASE_URL, SLOT);
        if (before) log(`[prefix-cache] slot before: tokens=${before.n_prompt_tokens}`);

        const res = await warm();
        const t = res.timings ?? {};
        log(
          `[prefix-cache] warm-up done: HTTP ${res.status} ` +
            `prompt_n=${t.prompt_n ?? "?"} cache_n=${t.cache_n ?? "?"}`,
        );

        // (c) persist the freshly built prefix (Phase 5)
        let saved = false;
        if (PERSIST) {
          const s = await saveSlotFile(BASE_URL, SLOT, fname);
          if (s.ok) {
            saved = true;
            log(`[prefix-cache] SAVED ${fname}: n_saved=${s.n_saved} (${s.timings?.save_ms?.toFixed(1) ?? "?"} ms)`);
            try {
              recordUse(fp, { ...meta, tokenCount: t.prompt_n ?? 0 });
            } catch (err) {
              log(`[prefix-cache] metadata update failed (fail-open): ${String(err)}`);
            }
          } else {
            log(`[prefix-cache] save failed (fail-open, HTTP ${s.status}${s.error ? `: ${s.error}` : ""})`);
          }
        }
        notify(
          `prefix-cache: COLD — warmed ${t.prompt_n ?? "?"} tok` +
            (PERSIST ? (saved ? ", checkpoint saved" : ", checkpoint save FAILED") : ""),
          "warning",
        );
      }

      if (BASE_URL) {
        try {
          const s = await getSlotStats(BASE_URL, SLOT);
          if (s) lastSeenPromptTokens = s.n_prompt_tokens;
        } catch {
          /* non-fatal */
        }
      }

      warmedFingerprint = fp;
      log(`[prefix-cache] normal request proceeding (source=${restored ? "restore" : "warm-up"})`);
    } catch (err) {
      log(`[prefix-cache] ERROR (fail-open): ${String(err)}`);
      notify(`prefix-cache: disabled this session (fail-open): ${String(err).slice(0, 120)}`, "error");
    }
  });

  // ---- Phase 6 diagnostics: observed reuse of the NORMAL request ----
  // after_provider_response carries no timings, and it fires before the
  // stream is consumed (prompt may still be processing). Slot counters
  // reflect the LAST prompt, so poll until n_prompt_tokens changes.
  pi.on("after_provider_response", async (event) => {
    if (!BASE_URL || !lastWasTarget || event.status >= 400) return;
    try {
      const deadline = Date.now() + 8000;
      for (;;) {
        await new Promise((r) => setTimeout(r, 500));
        const s = await getSlotStats(BASE_URL, SLOT);
        if (s && s.n_prompt_tokens !== lastSeenPromptTokens) {
          const pct =
            s.n_prompt_tokens > 0
              ? Math.round((100 * s.n_prompt_tokens_cache) / s.n_prompt_tokens)
              : 0;
          log(
            `[prefix-cache] observed normal request: prompt_n=${s.n_prompt_tokens} ` +
              `cache_n=${s.n_prompt_tokens_cache} (reuse ${pct}%)`,
          );
          lastSeenPromptTokens = s.n_prompt_tokens;
          return;
        }
        if (Date.now() > deadline) return;
      }
    } catch {
      /* fail-open: diagnostics only */
    }
  });
}