/**
 * favourite-index plugin — two-namespace tool presentation.
 *
 * Measured on a 145-tool catalog (scripts/model-eval/toolsim.py) to be the tool
 * presentation that most reduces wander and lifts first-pick for a wandering
 * model: expose a curated `favourite/<tool>` shortlist FIRST, the full
 * `all/<server>/<tool>` catalog second, plus a load-bearing instruction telling
 * the model to prefer favourites and fall back to `all/` only when nothing fits.
 * (The instruction is not optional — the same two namespaces without it scored
 * materially worse. See memory `feedback_llm_tool_scoping`.)
 *
 * Composes AFTER the gate/default plugin: it transforms the *ungated* upstream
 * catalog and no-ops while a session is still gated (only `begin_session`
 * visible). Gate/agent virtual tools (read_prompts, propose_prompt,
 * agent-<name> chat) pass through unchanged at the top level.
 *
 * Routing: the presented names are re-mapped back to their canonical
 * `server/tool` in `onToolCallBefore` (a name rewrite, then the chain
 * continues) so the normal upstream routing AND the content-pipeline
 * `onToolCallAfter` still run — unlike a virtual-tool alias, which would
 * short-circuit before content processing.
 */
import type { ProxyModelPlugin, PluginSessionContext } from '../plugin.js';
import type { ToolDefinition } from '../types.js';

/** Per-session state key holding the presented-name → canonical-name map. */
const RESOLVER_KEY = 'favourite-index:resolver';

/** The load-bearing "prefer favourite/" instruction (see module doc). */
export const FAVOURITE_INDEX_INSTRUCTION =
  'Tools are indexed in two namespaces. PREFER favourite/<tool> — a short ' +
  'curated list of the common tools that covers most tasks; reach for these ' +
  'first. Use the full catalog under all/<server>/<tool> only if nothing in ' +
  "favourites fits. (read_prompts gives this project's own guidance.)";

export interface FavouriteIndexConfig {
  /** Canonical `server/tool` names to surface as favourites, in display order. */
  favourites: string[];
  /** Cap on favourites presented (default: all provided). Enforces "not everything". */
  maxFavourites?: number;
}

function humanize(name: string): string {
  return name.replace(/_/g, ' ');
}

/** Split a canonical `server/tool` on the FIRST slash (server names have none). */
function splitServer(canonical: string): { server: string; short: string } {
  const i = canonical.indexOf('/');
  return i < 0
    ? { server: '', short: canonical }
    : { server: canonical.slice(0, i), short: canonical.slice(i + 1) };
}

interface Presentation {
  tools: ToolDefinition[];
  /** presented name → canonical `server/tool` */
  resolver: Map<string, string>;
}

/**
 * Reshape an upstream catalog into favourites-first + full all/ catalog, and
 * the presented→canonical routing map. Pure — shared by onToolsList and the
 * lazy rebuild in onToolCallBefore.
 */
function buildPresentation(
  upstreamTools: ToolDefinition[],
  favourites: string[],
  cap?: number,
): Presentation {
  const byName = new Map(upstreamTools.map((t) => [t.name, t]));
  const resolver = new Map<string, string>();
  const usedPresented = new Set<string>();
  const seenCanonical = new Set<string>();

  // Favourites first: configured order, only those present, capped, with
  // short-name collision disambiguation.
  const limit = cap ?? favourites.length;
  const favTools: ToolDefinition[] = [];
  for (const canonical of favourites) {
    if (favTools.length >= limit) break;
    if (seenCanonical.has(canonical)) continue;
    const t = byName.get(canonical);
    if (!t) continue; // stale/missing pin — skip (all/ still exposes it)
    const { server, short } = splitServer(canonical);
    let presented = `favourite/${short}`;
    if (usedPresented.has(presented)) presented = `favourite/${server}-${short}`;
    if (usedPresented.has(presented)) continue; // still collides (rare) — skip
    seenCanonical.add(canonical);
    usedPresented.add(presented);
    resolver.set(presented, canonical);
    favTools.push({
      name: presented,
      description: `${humanize(short)} — common (${server})`,
      ...(t.inputSchema !== undefined ? { inputSchema: t.inputSchema } : {}),
    });
  }

  // Then the full catalog under all/<server>/<tool>.
  const allTools: ToolDefinition[] = upstreamTools.map((t) => {
    const presented = `all/${t.name}`;
    resolver.set(presented, t.name);
    return {
      name: presented,
      ...(t.description !== undefined ? { description: t.description } : {}),
      ...(t.inputSchema !== undefined ? { inputSchema: t.inputSchema } : {}),
    };
  });

  return { tools: [...favTools, ...allTools], resolver };
}

export function createFavouriteIndexPlugin(config: FavouriteIndexConfig): ProxyModelPlugin {
  const favourites = config.favourites ?? [];
  const cap = config.maxFavourites && config.maxFavourites > 0 ? config.maxFavourites : undefined;

  return {
    name: 'favourite-index',
    description:
      'Two-namespace tool presentation: curated favourite/<tool> + full all/<server>/<tool>, prefer favourites.',

    async onInitialize() {
      return { instructions: FAVOURITE_INDEX_INSTRUCTION };
    },

    async onToolsList(tools: ToolDefinition[], ctx: PluginSessionContext): Promise<ToolDefinition[]> {
      // Distinguish upstream tools (in the discovered catalog) from virtual
      // tools (gate/agent) that must pass through untouched.
      const upstream = await ctx.discoverTools();
      const upstreamNames = new Set(upstream.map((t) => t.name));
      const upstreamTools = tools.filter((t) => upstreamNames.has(t.name));
      const passthrough = tools.filter((t) => !upstreamNames.has(t.name));

      // Gated / empty catalog → no-op (leave the list exactly as gate produced it).
      if (upstreamTools.length === 0) {
        ctx.state.delete(RESOLVER_KEY);
        return tools;
      }

      const { tools: presented, resolver } = buildPresentation(upstreamTools, favourites, cap);
      ctx.state.set(RESOLVER_KEY, resolver);
      return [...presented, ...passthrough];
    },

    async onToolCallBefore(toolName, _args, request, ctx) {
      if (!toolName.startsWith('favourite/') && !toolName.startsWith('all/')) return null;

      let resolver = ctx.state.get(RESOLVER_KEY) as Map<string, string> | undefined;
      let canonical = resolver?.get(toolName);

      // Resolver missing (a call arrived before tools/list this session): rebuild
      // it so favourite/ names — which, unlike all/, can't be prefix-stripped —
      // still route.
      if (canonical === undefined) {
        const upstream = await ctx.discoverTools();
        if (upstream.length > 0) {
          resolver = buildPresentation(upstream, favourites, cap).resolver;
          ctx.state.set(RESOLVER_KEY, resolver);
          canonical = resolver.get(toolName);
        }
        // all/ is always derivable by stripping the prefix.
        if (canonical === undefined && toolName.startsWith('all/')) {
          canonical = toolName.slice('all/'.length);
        }
      }

      if (canonical && canonical !== toolName) {
        // Rewrite presented name → canonical `server/tool`, then continue the
        // chain (return null) so normal routing + content-pipeline still run.
        const params = request.params as Record<string, unknown> | undefined;
        if (params) params['name'] = canonical;
      }
      return null;
    },
  };
}
