maestro

/**
 * @file Image analysis tool — loads local images, resizes with ffmpeg, returns base64 data URI.
 *
 * Uses ffmpeg (absolute path) to avoid PATH and native binary issues in the LM Studio sandbox.
 * Images are aggressively compressed to fit within LLM context windows.
 */

import { tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { readFile, stat, rm, mkdir } from "fs/promises";
import { extname, isAbsolute, resolve, join } from "path";
import { tmpdir } from "os";
import { randomBytes } from "crypto";
import { execFile } from "child_process";
import { promisify } from "util";
import { getFfmpegPath, getFfprobePath } from "./ffmpegPath";

const execFileAsync = promisify(execFile);

const SUPPORTED_EXTENSIONS = new Set([
  ".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp", ".tiff", ".tif",
]);

const MAX_BASE64_BYTES = 150_000;

/** Block system-sensitive paths. */
const BLOCKED_PREFIXES = ["/etc", "/var", "/usr", "/System", "/Library", "/private"];
function isSafePath(p: string): boolean {
  return !BLOCKED_PREFIXES.some(prefix => p.startsWith(prefix));
}

/**
 * Core resize logic — shared between single and batch analysis.
 */
export async function resizeImage(
  resolvedPath: string,
  maxDim: number,
): Promise<{ dataUri: string; originalWidth: number; originalHeight: number; bytes: number }> {
  const ffmpeg = await getFfmpegPath();
  const ffprobe = await getFfprobePath();
  const tmpDir = join(tmpdir(), `maestro-img-${randomBytes(6).toString("hex")}`);
  await mkdir(tmpDir, { recursive: true });
  const outputPath = join(tmpDir, "resized.jpg");

  try {
    let originalWidth = 0, originalHeight = 0;
    try {
      const { stdout } = await execFileAsync(ffprobe, [
        "-v", "quiet", "-print_format", "json",
        "-show_streams", "-select_streams", "v:0", resolvedPath,
      ]);
      const stream = JSON.parse(stdout)?.streams?.[0];
      originalWidth = stream?.width ?? 0;
      originalHeight = stream?.height ?? 0;
    } catch {}

    await execFileAsync(ffmpeg, [
      "-i", resolvedPath,
      "-vf", `scale='min(${maxDim},iw)':'min(${maxDim},ih)':force_original_aspect_ratio=decrease`,
      "-q:v", "8", "-y", outputPath,
    ]);

    let buf = await readFile(outputPath);
    if (buf.byteLength > MAX_BASE64_BYTES) {
      const smallerPath = join(tmpDir, "smaller.jpg");
      await execFileAsync(ffmpeg, [
        "-i", resolvedPath,
        "-vf", `scale='min(${Math.round(maxDim * 0.5)},iw)':'min(${Math.round(maxDim * 0.5)},ih)':force_original_aspect_ratio=decrease`,
        "-q:v", "10", "-y", smallerPath,
      ]);
      buf = await readFile(smallerPath);
    }

    return {
      dataUri: `data:image/jpeg;base64,${buf.toString("base64")}`,
      originalWidth, originalHeight, bytes: buf.byteLength,
    };
  } finally {
    await rm(tmpDir, { recursive: true, force: true }).catch(() => {});
  }
}

export function createImageAnalysisTool(ctl: ToolsProviderController, configMaxDim: number = 384): Tool {
  return tool({
    name: "analyze_image",
    description:
      "Load an image from a local file path, resize it to a small thumbnail, and return it as a base64 data URI for visual analysis. " +
      "Images are compressed to fit in the context window. " +
      "Supports JPEG, PNG, WebP, GIF, BMP, TIFF.",
    parameters: {
      file_path: z.string().describe("Absolute path to the image file."),
      max_dimension: z
        .number()
        .int()
        .min(64)
        .max(1280)
        .optional()
        .describe("Max width or height in pixels. Default: 512."),
    },
    implementation: async (
      { file_path, max_dimension }: { file_path: string; max_dimension?: number },
      { status },
    ) => {
      const maxDim = Math.min(max_dimension ?? configMaxDim, 1280);
      const resolvedPath = isAbsolute(file_path) ? file_path : resolve(file_path);

      if (!isSafePath(resolvedPath)) {
        return { error: `Access denied: '${resolvedPath}' is in a protected system directory.` };
      }
      const ext = extname(resolvedPath).toLowerCase();
      if (!SUPPORTED_EXTENSIONS.has(ext)) {
        return { error: `Unsupported image format '${ext}'. Supported: ${[...SUPPORTED_EXTENSIONS].join(", ")}` };
      }
      try {
        const fileStat = await stat(resolvedPath);
        if (!fileStat.isFile()) return { error: `Not a file: ${resolvedPath}` };
      } catch {
        return { error: `File not found: ${resolvedPath}` };
      }

      status("Resizing image...");
      try {
        const result = await resizeImage(resolvedPath, maxDim);
        status("Image ready");
        return {
          file_path: resolvedPath,
          original_width: result.originalWidth || undefined,
          original_height: result.originalHeight || undefined,
          resized_bytes: result.bytes,
          max_dimension_used: maxDim,
          data_uri: result.dataUri,
        };
      } catch (err: any) {
        return { error: `Failed to process image: ${err?.message || String(err)}`, file_path: resolvedPath };
      }
    },
  });
}

/**
 * Batch image analysis — analyzes all images in a directory.
 * Returns data URIs for each image (up to a limit) for visual inspection.
 */
export function createBatchImageAnalysisTool(ctl: ToolsProviderController, configMaxDim: number = 512): Tool {
  return tool({
    name: "analyze_images",
    description:
      "Analyze ALL images in a directory at once. Returns thumbnails for each image. " +
      "Much more efficient than calling analyze_image repeatedly. " +
      "Use this when you need to see multiple reference images or catalog a folder.",
    parameters: {
      directory: z.string().describe("Path to directory containing images."),
      max_dimension: z.number().int().min(64).max(1280).optional()
        .describe("Max width/height per image. Default: 256 for batch (smaller to fit more)."),
      limit: z.number().int().min(1).max(50).optional()
        .describe("Max number of images to process. Default: 20."),
    },
    implementation: async (
      { directory, max_dimension, limit: maxImages }: { directory: string; max_dimension?: number; limit?: number },
      { status },
    ) => {
      const maxDim = Math.min(max_dimension ?? 256, 1280);
      const imageLimit = maxImages ?? 20;
      const dirPath = isAbsolute(directory) ? directory : resolve(directory);

      if (!isSafePath(dirPath)) {
        return { error: `Access denied: '${dirPath}' is in a protected system directory.` };
      }

      let entries: string[];
      try {
        const { readdir } = await import("fs/promises");
        entries = await readdir(dirPath);
      } catch {
        return { error: `Cannot read directory: ${dirPath}` };
      }

      const imageFiles = entries
        .filter(name => SUPPORTED_EXTENSIONS.has(extname(name).toLowerCase()))
        .sort()
        .slice(0, imageLimit);

      if (imageFiles.length === 0) {
        return { error: `No supported images found in: ${dirPath}` };
      }

      status(`Processing ${imageFiles.length} images...`);

      // Process in parallel batches of 5 to avoid overwhelming ffmpeg
      const results: Array<{ name: string; data_uri: string; width: number; height: number } | { name: string; error: string }> = [];
      for (let i = 0; i < imageFiles.length; i += 5) {
        const batch = imageFiles.slice(i, i + 5);
        const batchResults = await Promise.all(
          batch.map(async (name) => {
            try {
              const fullPath = join(dirPath, name);
              const r = await resizeImage(fullPath, maxDim);
              return { name, data_uri: r.dataUri, width: r.originalWidth, height: r.originalHeight };
            } catch (err: any) {
              return { name, error: err?.message || String(err) };
            }
          }),
        );
        results.push(...batchResults);
        status(`Processed ${Math.min(i + 5, imageFiles.length)}/${imageFiles.length} images...`);
      }

      const successful = results.filter(r => "data_uri" in r).length;
      return {
        directory: dirPath,
        total_found: entries.filter(n => SUPPORTED_EXTENSIONS.has(extname(n).toLowerCase())).length,
        processed: results.length,
        successful,
        images: results,
      };
    },
  });
}

/**
 * Catalog images — returns metadata + short filenames for all images in a directory.
 * Lightweight alternative to analyze_images when you just need the inventory.
 */
export function createCatalogImagesTool(ctl: ToolsProviderController): Tool {
  return tool({
    name: "catalog_images",
    description:
      "List all images in a directory with metadata (name, dimensions, size). " +
      "Use this FIRST to understand available assets before analyzing specific ones. " +
      "No image data is returned — just the inventory.",
    parameters: {
      directory: z.string().describe("Path to directory containing images."),
    },
    implementation: async (
      { directory }: { directory: string },
      { status },
    ) => {
      const dirPath = isAbsolute(directory) ? directory : resolve(directory);
      if (!isSafePath(dirPath)) {
        return { error: `Access denied: '${dirPath}'` };
      }

      let entries: string[];
      try {
        const { readdir } = await import("fs/promises");
        entries = await readdir(dirPath);
      } catch {
        return { error: `Cannot read directory: ${dirPath}` };
      }

      const imageFiles = entries
        .filter(name => SUPPORTED_EXTENSIONS.has(extname(name).toLowerCase()))
        .sort();

      if (imageFiles.length === 0) {
        return { directory: dirPath, count: 0, images: [] };
      }

      status(`Cataloging ${imageFiles.length} images...`);

      const ffprobe = await getFfprobePath();
      const catalog = await Promise.all(
        imageFiles.map(async (name) => {
          const fullPath = join(dirPath, name);
          try {
            const fileStat = await stat(fullPath);
            let width = 0, height = 0;
            try {
              const { stdout } = await execFileAsync(ffprobe, [
                "-v", "quiet", "-print_format", "json",
                "-show_streams", "-select_streams", "v:0", fullPath,
              ]);
              const stream = JSON.parse(stdout)?.streams?.[0];
              width = stream?.width ?? 0;
              height = stream?.height ?? 0;
            } catch {}
            return {
              name,
              width, height,
              size_kb: Math.round(fileStat.size / 1024),
              ext: extname(name).toLowerCase(),
            };
          } catch {
            return { name, error: "unreadable" };
          }
        }),
      );

      return { directory: dirPath, count: catalog.length, images: catalog };
    },
  });
}