import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { registerPromptSegment } from "../../../vera-prompt-inspector/src/registry";
import { getAgentPath } from "../../src/shared/agent-paths";
import { emitDiagnostics } from "../../src/shared/diagnostics";
import { notify } from "../../src/shared/ui";
import { loadPromptRules } from "./rules-loader";

export default function promptRules(pi: ExtensionAPI) {
  // `initialized` distinguishes "not yet loaded" from "loaded but
  // found nothing". We set it to true BEFORE the load so that a load
  // error doesn't cause infinite re-attempts on every turn.
  //
  // Initialization runs from whichever lifecycle event fires first:
  // `session_start` on a cold boot, or `before_agent_start` after a
  // Pi extension reload (which replaces the closure but may not re-
  // emit `session_start`). Without this lazy path, a reloaded
  // extension would carry `assembledPrompt = null` for the rest of
  // the session and silently leave the system prompt at Pi's base.
  let initialized = false;
  let assembledPrompt: string | null = null;
  let layerCount = 0;
  let includeCount = 0;

  function ensureLoaded(ctx: ExtensionContext): void {
    if (initialized) return;
    initialized = true;
    const result = loadPromptRules(ctx.cwd);
    assembledPrompt = result.prompt;
    layerCount = result.layers.length;
    includeCount = result.includeCount;

    emitDiagnostics(ctx, result.diagnostics);

    if (!assembledPrompt) {
      notify(
        ctx,
        `prompt-rules: no rules found in ${ctx.cwd}/.pi/rules or ${getAgentPath("rules")} — system prompt unchanged`,
        "warning",
      );
      return;
    }

    const layerNames = result.layers.map((layer) => layer.name).join(", ");
    notify(
      ctx,
      `prompt-rules: loaded ${result.layers.length} layer(s) [${layerNames}] + ${result.includeCount} include(s)`,
      "info",
    );
  }

  pi.on("session_start", async (_event, ctx) => {
    ensureLoaded(ctx);
  });

  pi.on("before_agent_start", async (_event, ctx) => {
    ensureLoaded(ctx);
    registerPromptSegment({
      id: "soul-core",
      label: "soul + core rules",
      category: "rules",
      kind: "cached-per-session",
      text: assembledPrompt ?? "",
      source: "rules/SOUL.md + @include",
      details: [
        { label: "layers", value: String(layerCount) },
        { label: "includes", value: String(includeCount) },
      ],
    });
    // Do not modify event.systemPrompt here; vera-prompt-inspector
    // (loaded last) assembles the final systemPrompt from the registry.
    return undefined;
  });
}