find-image

import { tool, type Tool, type ToolCallContext, type ToolsProviderController } from "@lmstudio/sdk";
// @ts-ignore - zod/lib re-export chain breaks with NodeNext moduleResolution; runtime is fine
import { z } from "zod";
import path from "node:path";
import { resolveMediaQueries } from "../media/mediaResolver.js";
import { expandCachedImagePathsByContentIdentity, findCachedGenerationsByImagePaths, listCachedGenerationTags, updateGenerationTags } from "../indexer.js";
import { formatToolMetaBlock } from "../helpers/pluginMeta.js";

const TagImageParamsSchema = {
  action: z.enum(["list_tags", "show_tag", "add_tag", "remove_tag", "remove_all_tags"]).describe("list_tags lists all indexed tags; show_tag lists current tags; add_tag adds tags; remove_tag removes named tags; remove_all_tags clears every tag from each target."),
  target: z.union([z.array(z.string()), z.string()]).optional().describe("One or more indexed image references: vN, iN, pN, absolute image paths, or a concrete project:// thumbnail reference. Attachments (aN) and whole .sqlite3 project files cannot be tagged. Use an array for multiple targets."),
  tags: z.union([z.array(z.string()), z.string()]).optional().describe("Tags to add or remove. Use an array for multiple tags."),
};

interface TagImageArgs {
  action: "list_tags" | "show_tag" | "add_tag" | "remove_tag" | "remove_all_tags";
  target?: string[] | string;
  tags?: string[] | string;
}

function normalizeStringList(value: string[] | string | undefined): string[] {
  const rawValues = Array.isArray(value) ? value : value === undefined ? [] : [value];
  return rawValues.flatMap((rawValue) => {
    const value = rawValue.trim();
    if (value.startsWith("[") && value.endsWith("]")) {
      try {
        const parsed = JSON.parse(value);
        if (Array.isArray(parsed)) return parsed.map((item) => String(item ?? "").trim()).filter(Boolean);
      } catch {
        // Keep the original string for the normal parser.
      }
    }
    return [value];
  }).filter(Boolean);
}

function normalizeTags(value: TagImageArgs["tags"]): string[] {
  const rawTags = normalizeStringList(value).flatMap((tag) => tag.split(/[\s,]+/));
  const tagsByNormalizedValue = new Map<string, string>();
  for (const rawTag of rawTags) {
    const tag = rawTag.trim();
    if (tag) tagsByNormalizedValue.set(tag.toLocaleLowerCase(), tag);
  }
  return [...tagsByNormalizedValue.values()];
}

function normalizeTargetNotationToken(raw: unknown): string {
  let target = String(raw ?? "").trim();
  if (!target || path.isAbsolute(target)) return target;
  target = target.replace(/^[\[{(]+/, "").replace(/[\]})]+$/, "");
  target = target.replace(/^['"`]+|['"`]+$/g, "");
  target = target.replace(/[;:.!?]+$/, "");
  if (target.startsWith("project://")) return target;
  return target.trim().toLowerCase();
}

function normalizeTargets(value: TagImageArgs["target"]): string[] {
  const rawTargets = normalizeStringList(value);
  const targets = rawTargets.flatMap((rawTarget) => {
    const target = normalizeTargetNotationToken(rawTarget);
    return /^(?:[avip]\d+)(?:[\s,]+[avip]\d+)+$/i.test(target)
      ? target.split(/[\s,]+/)
      : [target];
  }).filter(Boolean);
  return [...new Set(targets)];
}

async function resolveTargetImagePaths(ctl: ToolsProviderController, target: string): Promise<string[]> {
  if (/^\s*a\d+\s*$/i.test(target)) {
    throw new Error(`Attachments cannot be tagged: ${target}.`);
  }
  if (/^\s*[vip]\d+\s*$/i.test(target)) {
    const workingDirectory = ctl.getWorkingDirectory();
    if (typeof workingDirectory !== "string" || !workingDirectory.trim()) return [];
    const resolved = await resolveMediaQueries(workingDirectory, target);
    return resolved.mediaQueries.map((media) => media.filePath).filter((imagePath): imagePath is string => !!imagePath);
  }
  const candidatePath = target.trim().replace(/^['"]|['"]$/g, "");
  if (candidatePath.startsWith("project://")) {
    if (!/^project:\/\/.+\.sqlite3#\d+$/i.test(candidatePath)) {
      throw new Error(`A project target must name one thumbnail: ${target}.`);
    }
    return [candidatePath];
  }
  if (path.extname(candidatePath).toLowerCase() === ".sqlite3") {
    throw new Error(`Project files cannot be tagged: ${target}.`);
  }
  return path.isAbsolute(candidatePath) ? [candidatePath] : [];
}

export function createTagImageTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "tag_image",
    description: `Manage persistent tags for indexed images.
  - action: list_tags, show_tag, add_tag, remove_tag, or remove_all_tags
  - target: one or more vN, iN, pN, absolute image paths, or concrete project:// thumbnail references that still resolve to indexed images; attachments (aN) and whole .sqlite3 project files cannot be tagged
- tags: one or more tags; use an array for multiple tags

Examples:
- { "action": "add_tag", "target": "p1", "tags": ["favorite"] }
- { "action": "remove_tag", "target": ["p1", "v2", "/absolute/path/image.png"], "tags": ["favorite", "reviewed"] }

${formatToolMetaBlock()}`,
    parameters: TagImageParamsSchema,
    implementation: async (args: TagImageArgs, ctx: ToolCallContext) => {
      try {
        if (args.action === "list_tags") {
          const tags = listCachedGenerationTags();
          ctx.status("Tags loaded.");
          return JSON.stringify({ type: "tag-image-result", action: args.action, tags, images: [] }, null, 2);
        }

        const targets = normalizeTargets(args.target);
        if (targets.length === 0) throw new Error("tag_image requires at least one target image.");
        const tags = normalizeTags(args.tags);
        if (args.action !== "show_tag" && args.action !== "remove_all_tags" && tags.length === 0) {
          throw new Error(`${args.action} requires at least one tag.`);
        }
        ctx.status("Resolving image...");
        const resolvedImagePaths = [...new Set((await Promise.all(targets.map((target) => resolveTargetImagePaths(ctl, target)))).flat())];
        const imagePaths = await expandCachedImagePathsByContentIdentity(resolvedImagePaths);
        if (imagePaths.length === 0) throw new Error(`No targets resolved: ${targets.join(", ")}.`);

        const matched = findCachedGenerationsByImagePaths(imagePaths);
        if (matched.length === 0) throw new Error(`Targets are not indexed: ${targets.join(", ")}.`);

        const updated = args.action === "show_tag" ? matched : updateGenerationTags(imagePaths, args.action, tags);
        const images = updated.map((generation) => ({ imagePaths: generation.imagePaths ?? [], tags: generation.tags ?? [] }));
        ctx.status(args.action === "show_tag" ? "Tags loaded." : "Tags updated.");
        return JSON.stringify({ type: "tag-image-result", action: args.action, targets, tags: images[0]?.tags ?? [], images }, null, 2);
      } catch (error) {
        return JSON.stringify({ type: "tag-image-result", action: args.action, targets: normalizeTargets(args.target), error: String((error as Error).message ?? error), images: [] }, null, 2);
      }
    },
  });
}