/**
 * pi-tool-search — model-built tool groups with bounded dynamic schema loading.
 *
 * Imported from https://github.com/tuansondinh/pi-tool-search and maintained
 * locally from snapshot ddfb23646fd3957b791214de278e23aa393c9b13 (v0.3.6).
 */
import { randomUUID } from "node:crypto";
import { unlinkSync } from "node:fs";
import { join } from "node:path";

import { getAgentDir, type ExtensionAPI, type ExtensionContext, type ToolInfo } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

import {
  buildCatalogModelInput,
  buildManifestHash,
  createFallbackCatalog,
  parseGeneratedCatalog,
  rankGroups,
  readCachedCatalog,
  selectConfidentGroup,
  writeCachedCatalog,
  type CatalogConstraints,
  type GroupCard,
  type ToolCatalog,
} from "./catalog.ts";
import {
  catalogPreservesBundleAssignments,
  createBundleCatalog,
  type BundleCatalogResult,
} from "./bundle-groups.ts";
import { ensureToolSearchDefaults, readToolSearchConfig } from "./config.ts";

const TOOL_SEARCH_NAME = "tool_search";
const POWERSHELL_TOOL_NAME = "powershell";
const MAX_BATCH_GROUPS = 5;
const CORE_TOOLS = ["read", "write", "edit", "bash", "grep", "find"];

type CatalogSource = "bundle" | "cache" | "fallback" | "hybrid" | "model";
type ModelUsage = Awaited<ReturnType<ExtensionContext["modelRegistry"]["complete"]>>["usage"];
type CatalogGenerationResult = { notice?: string; usage?: ModelUsage };

export function platformToolPolicy<T extends { name: string }>(
  definitions: readonly T[],
  platform: NodeJS.Platform,
): { tools: T[]; coreTools: string[] } {
  const tools = platform === "win32"
    ? [...definitions]
    : definitions.filter((tool) => tool.name !== POWERSHELL_TOOL_NAME);
  const coreTools = platform === "win32"
    ? [...CORE_TOOLS, POWERSHELL_TOOL_NAME]
    : [...CORE_TOOLS];
  return { tools, coreTools };
}

function unique(values: Iterable<string>): string[] {
  return [...new Set(values)];
}

function catalogPrompt(toolsJson: string, maxToolsPerGroup: number, fixedGroupsJson: string): string {
  return [
    "Build a compact retrieval catalog for the provided Pi tools.",
    "The catalog is metadata only: never rename tools or invent parameters.",
    "Group tools that are commonly needed together in one workflow, but keep providers or security-sensitive administration separate when that improves routing.",
    `Every tool must appear exactly once and every group must contain 1-${maxToolsPerGroup} tools.`,
    "Tools listed in <fixed-groups> must keep exactly those primaryGroup ids; these are curated bundle assignments.",
    "Summaries must preserve capability, important boundaries, and when not to use a tool.",
    "Keywords should include common English terms and useful Chinese equivalents when applicable.",
    "Return JSON only with this exact shape:",
    '{"groups":[{"id":"kebab-case","title":"...","summary":"...","useWhen":["..."],"avoidWhen":["..."],"tools":["exact_tool_name"]}],"tools":[{"name":"exact_tool_name","summary":"...","useWhen":["..."],"avoidWhen":["..."],"keywords":["..."],"primaryGroup":"kebab-case"}]}',
    "Do not include markdown fences or explanatory text.",
    "",
    "<fixed-groups>",
    fixedGroupsJson,
    "</fixed-groups>",
    "",
    "<tools>",
    toolsJson,
    "</tools>",
  ].join("\n");
}

function responseText(content: Array<{ type: string; text?: string }>): string {
  return content
    .filter((item): item is { type: string; text: string } => item.type === "text" && typeof item.text === "string")
    .map((item) => item.text)
    .join("\n");
}

function groupDescription(catalog: ToolCatalog, source: CatalogSource): string {
  const sourceLabel =
    source === "bundle"
      ? "precomputed my-pi bundle catalog; no model generation required"
      : source === "hybrid"
        ? "my-pi bundle catalog plus deterministic groups for unrecognized tools; the first search may enrich the unknown tools"
        : source === "fallback"
          ? "deterministic fallback"
          : `${source} catalog`;
  const groups = catalog.groups
    .map((group) => {
      const when = group.useWhen.slice(0, 2).join("; ");
      const detail = [group.summary, when].filter(Boolean).join(" Use when: ").slice(0, 360);
      return `  ${group.id}: ${detail}`;
    })
    .join("\n");
  return [
    "Activate one or more complete tool groups for the current task. Prefer exact group ids from this catalog; use query only when no id clearly matches.",
    `Catalog source: ${sourceLabel}.`,
    "Available groups:",
    groups || "  (no hidden tool groups)",
  ].join("\n");
}

export default function toolSearchExtension(pi: ExtensionAPI): void {
  const agentDir = getAgentDir();
  const cachePath = join(agentDir, "tool-search", "catalog-v1.json");
  const defaultResult = ensureToolSearchDefaults(agentDir);
  if (defaultResult === "skipped-invalid") {
    console.warn("my-pi: skipped pi-tool-search defaults because settings.json or toolSearch is invalid");
  }

  let config = readToolSearchConfig(agentDir);
  let tools: ToolInfo[] = [];
  let catalog: ToolCatalog = createFallbackCatalog([], "", {
    maxToolsPerGroup: Math.min(config.maxToolsPerGroup, config.maxDynamicTools),
    groupOverrides: {},
  });
  let catalogSource: CatalogSource = "fallback";
  let manifestHash = "";
  let policySignature = "";
  let attemptedGenerationHash: string | undefined;
  let clock = 0;
  const activeGroups = new Map<string, number>();
  const pinnedTools = new Set<string>();
  let bundleState: BundleCatalogResult = { catalog, coveredNames: new Set(), unknownTools: [] };
  let lastEviction = "none";

  function constraints(): CatalogConstraints {
    return {
      maxToolsPerGroup: Math.min(config.maxToolsPerGroup, config.maxDynamicTools),
      groupOverrides: config.groupOverrides,
    };
  }

  function installBaseCatalog(useCache: boolean): void {
    bundleState = createBundleCatalog(tools, manifestHash, constraints());
    if (bundleState.unknownTools.length === 0) {
      catalog = bundleState.catalog;
      catalogSource = "bundle";
      return;
    }
    const cached = useCache ? readCachedCatalog(cachePath, tools, manifestHash, constraints()) : undefined;
    if (cached && catalogPreservesBundleAssignments(cached, bundleState.catalog, bundleState.coveredNames)) {
      catalog = cached;
      catalogSource = "cache";
      return;
    }
    catalog = bundleState.catalog;
    catalogSource = "hybrid";
  }

  function fixedBundleGroupsJson(): string {
    return JSON.stringify(
      bundleState.catalog.groups
        .map((group) => ({
          id: group.id,
          tools: group.tools.filter((name) => bundleState.coveredNames.has(name)),
        }))
        .filter((group) => group.tools.length > 0),
    );
  }

  function groupById(id: string): GroupCard | undefined {
    return catalog.groups.find((group) => group.id === id);
  }

  function dynamicToolCount(groupIds: Iterable<string> = activeGroups.keys()): number {
    return unique([...groupIds].flatMap((id) => groupById(id)?.tools ?? [])).length;
  }

  function leastRecentlyUsedGroup(
    groups: Map<string, number> = activeGroups,
    protectedGroups: ReadonlySet<string> = new Set(),
  ): string | undefined {
    return [...groups]
      .filter(([id]) => !protectedGroups.has(id))
      .sort((left, right) => left[1] - right[1] || left[0].localeCompare(right[0]))[0]?.[0];
  }

  function currentStatusText(): string {
    const lru = [...activeGroups]
      .sort((left, right) => left[1] - right[1] || left[0].localeCompare(right[0]))
      .map(([id]) => id);
    return [
      "Tool Search status",
      `Catalog: ${catalogSource}`,
      `Pinned tools: ${pinnedTools.size}`,
      `Dynamic groups: ${activeGroups.size}/${config.maxActiveGroups}`,
      `Dynamic tools: ${dynamicToolCount()}/${config.maxDynamicTools}`,
      `LRU oldest → newest: ${lru.join(" → ") || "none"}`,
      `Last eviction: ${lastEviction}`,
    ].join("\n");
  }

  function applyActiveTools(): void {
    const groupTools = [...activeGroups.keys()].flatMap((id) => groupById(id)?.tools ?? []);
    pi.setActiveTools(unique([TOOL_SEARCH_NAME, ...pinnedTools, ...groupTools]));
  }

  function updateStatus(ctx: Pick<ExtensionContext, "ui">): void {
    const activeToolCount = unique([TOOL_SEARCH_NAME, ...pinnedTools, ...[...activeGroups.keys()].flatMap((id) => groupById(id)?.tools ?? [])]).length;
    ctx.ui.setStatus(
      "tool-search",
      config.showToolSearchFooterStatus
        ? `${activeToolCount} / ${tools.length + pinnedTools.size + 1} tools · ${activeGroups.size} / ${config.maxActiveGroups} groups`
        : undefined,
    );
  }

  function registerToolSearch(): void {
    pi.registerTool({
      name: TOOL_SEARCH_NAME,
      label: "Tool Search",
      description: groupDescription(catalog, catalogSource),
      promptSnippet: `Activate relevant tool groups on demand; up to ${config.maxActiveGroups} dynamic groups can remain active. Use groups for workflows that need several groups together`,
      parameters: Type.Object({
        group: Type.Optional(Type.String({ description: "Exact group id from the tool_search catalog" })),
        groups: Type.Optional(
          Type.Array(Type.String({ description: "Exact group id from the tool_search catalog" }), {
            description: `Exact group ids to activate atomically (maximum ${MAX_BATCH_GROUPS})`,
            minItems: 1,
            maxItems: MAX_BATCH_GROUPS,
          }),
        ),
        query: Type.Optional(Type.String({ description: "Natural-language task used to rank groups when an exact id is unclear" })),
      }),
      async execute(_toolCallId, params, signal, _onUpdate, ctx) {
        const singleGroup = typeof params.group === "string" ? params.group.trim() : "";
        const batchGroups = Array.isArray(params.groups) ? params.groups : [];
        const query = typeof params.query === "string" ? params.query.trim() : "";
        const hasSingle = singleGroup.length > 0;
        const hasBatch = batchGroups.length > 0;
        const hasQuery = query.length > 0;
        const modes = Number(hasSingle) + Number(hasBatch) + Number(hasQuery);
        const lines: string[] = [];
        const emptyDetails = (candidates: Array<{ id: string; score: number }> = []) => ({
          loadedGroup: undefined,
          loadedGroups: [] as string[],
          evictedGroups: [] as string[],
          activeGroups: [...activeGroups.keys()],
          candidates,
          catalogSource,
        });

        if (modes !== 1) {
          lines.push("Provide exactly one of group, groups, or query.");
          return { content: [{ type: "text", text: lines.join("\n") }], details: emptyDetails() };
        }
        if (batchGroups.length > MAX_BATCH_GROUPS) {
          lines.push(`groups accepts at most ${MAX_BATCH_GROUPS} group ids.`);
          return { content: [{ type: "text", text: lines.join("\n") }], details: emptyDetails() };
        }

        const requestedIds = unique(
          (hasBatch ? batchGroups : hasSingle ? [singleGroup] : [])
            .map((id) => id.trim().toLowerCase())
            .filter(Boolean),
        );
        let selectedGroups = requestedIds.map(groupById).filter((group): group is GroupCard => group !== undefined);
        const needsCatalog = hasQuery || selectedGroups.length !== requestedIds.length;
        const generation: CatalogGenerationResult = needsCatalog ? await ensureModelCatalog(ctx, signal) : {};
        if (generation.notice) lines.push(generation.notice);

        selectedGroups = requestedIds.map(groupById).filter((group): group is GroupCard => group !== undefined);
        const ranked = hasQuery ? rankGroups(catalog, query).slice(0, 3) : [];
        if (hasQuery) {
          const selected = selectConfidentGroup(ranked, query);
          if (selected) selectedGroups = [selected];
        }

        const missingIds = requestedIds.filter((id) => !groupById(id));
        if (selectedGroups.length === 0 || missingIds.length > 0) {
          const candidates = ranked.length > 0 ? ranked : catalog.groups.slice(0, 5).map((group) => ({ group, score: 0 }));
          lines.push(
            missingIds.length > 0
              ? `Unknown groups: ${missingIds.join(", ")}`
              : hasQuery
                ? "No group was activated because the query match was weak or ambiguous. Use an exact group id from the candidates."
                : "No group was activated. Provide an exact group id from the catalog.",
            `Candidates: ${candidates.map(({ group }) => group.id).join(", ") || "none"}`,
          );
          return {
            content: [{ type: "text", text: lines.join("\n") }],
            details: emptyDetails(candidates.map(({ group, score }) => ({ id: group.id, score }))),
            ...(generation.usage ? { usage: generation.usage } : {}),
          };
        }

        const selectedIds = selectedGroups.map((group) => group.id);
        const selectedTools = unique(selectedGroups.flatMap((group) => group.tools));
        if (selectedIds.length > config.maxActiveGroups || selectedTools.length > config.maxDynamicTools) {
          lines.push(
            `Requested workflow needs ${selectedIds.length} groups and ${selectedTools.length} tools, exceeding the configured limits of ${config.maxActiveGroups} groups and ${config.maxDynamicTools} tools.`,
          );
          return {
            content: [{ type: "text", text: lines.join("\n") }],
            details: emptyDetails(),
            ...(generation.usage ? { usage: generation.usage } : {}),
          };
        }

        const requestedSet = new Set(selectedIds);
        const nextActive = new Map(activeGroups);
        for (const id of selectedIds) if (!nextActive.has(id)) nextActive.set(id, Number.MAX_SAFE_INTEGER);
        const evictedGroups: string[] = [];
        const evictionReasons = new Set<string>();
        while (nextActive.size > config.maxActiveGroups || dynamicToolCount(nextActive.keys()) > config.maxDynamicTools) {
          if (nextActive.size > config.maxActiveGroups) evictionReasons.add(`group cap ${config.maxActiveGroups}`);
          if (dynamicToolCount(nextActive.keys()) > config.maxDynamicTools) evictionReasons.add(`tool cap ${config.maxDynamicTools}`);
          const evicted = leastRecentlyUsedGroup(nextActive, requestedSet);
          if (!evicted) break;
          nextActive.delete(evicted);
          evictedGroups.push(evicted);
        }

        const previousActive = new Map(activeGroups);
        const previousClock = clock;
        const previousLastEviction = lastEviction;
        activeGroups.clear();
        for (const [id, recency] of nextActive) activeGroups.set(id, recency);
        for (const id of selectedIds) activeGroups.set(id, ++clock);
        if (evictedGroups.length > 0) {
          lastEviction = `${evictedGroups.join(", ")} (${[...evictionReasons].join("; ")})`;
        }
        try {
          applyActiveTools();
          registerToolSearch();
          updateStatus(ctx);
        } catch (error) {
          activeGroups.clear();
          for (const [id, recency] of previousActive) activeGroups.set(id, recency);
          clock = previousClock;
          lastEviction = previousLastEviction;
          try {
            applyActiveTools();
            registerToolSearch();
            updateStatus(ctx);
          } catch {
            // Preserve the original host activation failure while retaining rolled-back internal state.
          }
          throw error;
        }

        if (selectedGroups.length === 1) {
          lines.push(`Loaded group: ${selectedGroups[0]?.id} (${selectedTools.join(", ")})`);
        } else {
          lines.push(`Loaded groups: ${selectedIds.join(", ")} (${selectedTools.length} unique tools)`);
        }
        if (evictedGroups.length > 0) lines.push(`Evicted least-recently-used: ${evictedGroups.join(", ")}`);
        lines.push(`Active groups: ${[...activeGroups.keys()].join(", ")}`);
        return {
          content: [{ type: "text", text: lines.join("\n") }],
          details: {
            loadedGroup: selectedIds.at(-1),
            loadedGroups: selectedIds,
            evictedGroups,
            activeGroups: [...activeGroups.keys()],
            candidates: ranked.map(({ group, score }) => ({ id: group.id, score })),
            catalogSource,
          },
          ...(generation.usage ? { usage: generation.usage } : {}),
        };
      },
    });
  }

  async function ensureModelCatalog(
    ctx: ExtensionContext,
    signal: AbortSignal | undefined,
  ): Promise<{ notice?: string; usage?: ModelUsage }> {
    if (catalogSource !== "hybrid" || attemptedGenerationHash === manifestHash || tools.length === 0) return {};
    attemptedGenerationHash = manifestHash;
    if (!ctx.model || !ctx.modelRegistry.hasConfiguredAuth(ctx.model)) {
      return { notice: "Using the precomputed bundle catalog plus deterministic groups for unrecognized tools because the current model is unavailable or unauthenticated." };
    }
    let generationUsage: ModelUsage | undefined;

    try {
      const response = await ctx.modelRegistry.complete(
        ctx.model,
        {
          messages: [
            {
              role: "user",
              content: [
                {
                  type: "text",
                  text: catalogPrompt(
                    buildCatalogModelInput(tools),
                    constraints().maxToolsPerGroup,
                    fixedBundleGroupsJson(),
                  ),
                },
              ],
              timestamp: Date.now(),
            },
          ],
        },
        {
          signal,
          reasoningEffort: "low",
          cacheRetention: "none",
          sessionId: randomUUID(),
          maxTokens: 12_000,
        },
      );
      generationUsage = response.usage;
      catalog = parseGeneratedCatalog(
        responseText(response.content),
        tools,
        manifestHash,
        constraints(),
        `${ctx.model.provider}/${ctx.model.id}`,
      );
      if (!catalogPreservesBundleAssignments(catalog, bundleState.catalog, bundleState.coveredNames)) {
        throw new Error("generated catalog changed one or more curated my-pi bundle assignments");
      }
      catalogSource = "model";
      try {
        writeCachedCatalog(cachePath, catalog);
      } catch (error) {
        console.warn(`my-pi: could not cache pi-tool-search catalog: ${error instanceof Error ? error.message : String(error)}`);
      }
      registerToolSearch();
      return {
        notice: `Generated and cached ${catalog.groups.length} groups while preserving the precomputed my-pi assignments.`,
        usage: generationUsage,
      };
    } catch (error) {
      return {
        notice: `Model enrichment failed; using the precomputed bundle catalog plus deterministic unknown-tool groups (${error instanceof Error ? error.message : String(error)}).`,
        ...(generationUsage ? { usage: generationUsage } : {}),
      };
    }
  }

  function refreshState(ctx: Pick<ExtensionContext, "ui">, forceReset: boolean): void {
    const nextConfig = readToolSearchConfig(agentDir);
    const registeredTools = pi.getAllTools().filter((tool) => tool.name !== TOOL_SEARCH_NAME);
    const platformPolicy = platformToolPolicy(registeredTools, process.platform);
    const allTools = platformPolicy.tools;
    const availableNames = new Set(allTools.map((tool) => tool.name));
    const nextPinned = new Set(
      [...platformPolicy.coreTools, ...nextConfig.alwaysEnabled].filter((name) => availableNames.has(name)),
    );
    const hiddenTools = allTools.filter((tool) => !nextPinned.has(tool.name));
    const nextConstraints = {
      maxToolsPerGroup: Math.min(nextConfig.maxToolsPerGroup, nextConfig.maxDynamicTools),
      groupOverrides: nextConfig.groupOverrides,
    };
    const nextHash = buildManifestHash(hiddenTools, nextConstraints);
    const nextPolicySignature = JSON.stringify({
      alwaysEnabled: [...nextPinned].sort(),
      maxActiveGroups: nextConfig.maxActiveGroups,
      maxDynamicTools: nextConfig.maxDynamicTools,
      showToolSearchFooterStatus: nextConfig.showToolSearchFooterStatus,
    });
    const catalogChanged = forceReset || nextHash !== manifestHash;
    const policyChanged = nextPolicySignature !== policySignature;

    config = nextConfig;
    tools = hiddenTools;
    pinnedTools.clear();
    for (const name of nextPinned) pinnedTools.add(name);
    manifestHash = nextHash;
    policySignature = nextPolicySignature;

    if (catalogChanged) {
      activeGroups.clear();
      clock = 0;
      lastEviction = "none";
      attemptedGenerationHash = undefined;
      installBaseCatalog(true);
    }

    const refreshEvictions: string[] = [];
    const refreshReasons = new Set<string>();
    while (activeGroups.size > config.maxActiveGroups || dynamicToolCount() > config.maxDynamicTools) {
      if (activeGroups.size > config.maxActiveGroups) refreshReasons.add(`group cap ${config.maxActiveGroups}`);
      if (dynamicToolCount() > config.maxDynamicTools) refreshReasons.add(`tool cap ${config.maxDynamicTools}`);
      const evicted = leastRecentlyUsedGroup();
      if (!evicted) break;
      activeGroups.delete(evicted);
      refreshEvictions.push(evicted);
    }
    const capacityChanged = refreshEvictions.length > 0;
    if (capacityChanged) lastEviction = `${refreshEvictions.join(", ")} (${[...refreshReasons].join("; ")}; settings refresh)`;
    if (catalogChanged || policyChanged || capacityChanged) {
      registerToolSearch();
      applyActiveTools();
    }
    updateStatus(ctx);
  }

  pi.registerCommand("tool-search-status", {
    description: "Show Tool Search catalog, capacity, active-group LRU order, and the latest eviction reason",
    handler: async (_args, ctx) => {
      if (ctx.hasUI) ctx.ui.notify(currentStatusText(), "info");
    },
  });

  pi.registerCommand("tool-search-rebuild", {
    description: "Invalidate the generated tool-group catalog; rebuild lazily on the next tool_search call",
    handler: async (_args, ctx) => {
      try {
        unlinkSync(cachePath);
      } catch (error) {
        if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
      }
      activeGroups.clear();
      attemptedGenerationHash = undefined;
      lastEviction = "none";
      installBaseCatalog(false);
      registerToolSearch();
      applyActiveTools();
      updateStatus(ctx);
      if (ctx.hasUI) {
        const message =
          bundleState.unknownTools.length === 0
            ? "Restored the precomputed my-pi tool-group catalog; no model rebuild is needed"
            : "Restored my-pi groups; unrecognized tools may be enriched on the next tool_search call";
        ctx.ui.notify(message, "info");
      }
    },
  });

  pi.on("session_start", (_event, ctx) => {
    refreshState(ctx, true);
  });

  pi.on("turn_start", (_event, ctx) => {
    refreshState(ctx, false);
  });

  pi.on("tool_execution_start", (event) => {
    const group = catalog.tools.find((tool) => tool.name === event.toolName)?.primaryGroup;
    if (group && activeGroups.has(group)) activeGroups.set(group, ++clock);
  });
}