import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { registerPromptSegment } from "../../../vera-prompt-inspector/src/registry";
import { emitDiagnostics } from "../../src/shared/diagnostics";
import { notify } from "../../src/shared/ui";
import { loadConditionalRules, type ConditionalRule } from "./loader";
import { matchesAny } from "./matcher";
import { resetTracker, recordAccess, getAccessedPaths } from "./tracker";
import { prescanWorkingTree } from "./prescan";

export default function conditionalRules(pi: ExtensionAPI) {
  // Lazy-init pattern: initialization runs from whichever lifecycle
  // event fires first. After a Pi extension reload the closure is
  // recreated with empty state but `session_start` is not re-emitted,
  // so we also bootstrap from `before_agent_start` / `tool_execution_end`.
  // `initialized` flips to true BEFORE loading so a loader error
  // cannot cause infinite retry.
  let initialized = false;
  let loadedRules: ConditionalRule[] = [];
  let manualActivated = new Set<string>();

  function ensureLoaded(ctx: ExtensionContext): void {
    if (initialized) return;
    initialized = true;
    resetTracker();
    manualActivated = new Set();
    const result = loadConditionalRules(ctx.cwd);
    loadedRules = result.rules;
    emitDiagnostics(ctx, result.diagnostics);

    const alwaysRules = loadedRules.filter((rule) => rule.when === "always-if-present");
    if (alwaysRules.length > 0) {
      const found = prescanWorkingTree(ctx.cwd, alwaysRules);
      for (const path of found) recordAccess(ctx.cwd, path);
    }

    if (loadedRules.length > 0) {
      notify(ctx, `conditional-rules: ${loadedRules.length} rule(s) loaded`, "info");
    }
  }

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

  pi.on("tool_execution_end", async (event, ctx) => {
    ensureLoaded(ctx);
    const tool = event?.tool ?? event?.toolName;
    if (tool !== "read" && tool !== "edit" && tool !== "write") return;
    const path = event?.args?.path;
    if (typeof path === "string") recordAccess(ctx.cwd, path);
  });

  pi.on("before_agent_start", async (_event, ctx) => {
    ensureLoaded(ctx);
    const accessed = getAccessedPaths();
    const matched = loadedRules.length === 0 ? [] : loadedRules.filter((rule) => {
      if (rule.when === "manual") return manualActivated.has(rule.name);
      return accessed.some((path) => matchesAny(path, rule.globs));
    });
    matched.sort((a, b) => a.priority - b.priority);
    const block = matched.length > 0 ? formatBlock(matched) : "";
    registerPromptSegment({
      id: "conditional",
      label: "conditional rules",
      category: "rules",
      kind: "per-turn",
      text: block,
      details: [
        { label: "loaded", value: String(loadedRules.length) },
        { label: "matched", value: String(matched.length) },
      ],
    });
    return undefined;
  });
}

function formatBlock(rules: ConditionalRule[]): string {
  const sections = rules.map((rule) => `### ${rule.name} (${rule.layer})\n\n${rule.body.trim()}`);
  return `## Conditional Rules (activated by working set)\n\n${sections.join("\n\n")}`;
}