draw-things-chat

/**
 * Tool-Result Harvester
 *
 * Unified harvesting of media from tool results.
 * Uses toolParams configs to determine extraction, preview, and injection behavior.
 *
 * Phase 3 of toolParams integration.
 */

import type { GeneratorController } from "@lmstudio/sdk";
import {
  getToolExtractor,
  getMediaTypeFullConfig,
  resolveActions,
  resolveInjection,
  type ToolResultCandidate,
  type MediaTypeFullConfig,
  type HarvestingActions,
  type InjectionConfig,
  type ChatMediaState,
  appendPictures,
  appendImages,
  generatePreview,
  previewFilenameFrom,
  isProjectUri,
  resolveProjectUri,
} from "../core-bundle.mjs";
// Note: trackInjection requires chatId, which this module doesn't have access to.
// Injection tracking is handled at the orchestrator level instead.
import fs from "fs";
import path from "path";
import crypto from "crypto";

/* ═══════════════════════════════════════════════════════════════════════════
 * TYPES
 * ═══════════════════════════════════════════════════════════════════════════ */

export interface ToolCallInfo {
  toolCallId?: string;
  toolName: string;
  pluginId?: string;
}

export interface HarvestContext {
  chatWd: string;
  state: ChatMediaState;
  toggles: Record<string, boolean>;
  debug?: boolean;
}

export interface HarvestResult {
  mediaType: "picture" | "image" | "variant" | "attachment";
  candidates: ToolResultCandidate[];
  actions: HarvestingActions;
  injection: InjectionConfig;
  config: MediaTypeFullConfig;
  stateChanged: boolean;
  markdown?: string;
}

export interface PreviewOptions {
  maxDim: number;
  quality: number;
}

/* ═══════════════════════════════════════════════════════════════════════════
 * HELPERS
 * ═══════════════════════════════════════════════════════════════════════════ */

function shortHexSha256(input: string, chars = 12): string {
  const h = crypto.createHash("sha256").update(input).digest("hex");
  return h.slice(0, Math.max(4, Math.min(64, chars)));
}

function canonicalizeExternalUrl(u: string): string {
  let s = String(u || "").trim();
  s = s.replace(/[)\],.;]+$/g, "");
  s = s.replace(/^"(.+)"$/g, "$1");
  s = s.replace(/^'(.+)'$/g, "$1");
  return s.trim();
}

/* ═══════════════════════════════════════════════════════════════════════════
 * MAIN HARVEST FUNCTION
 * ═══════════════════════════════════════════════════════════════════════════ */

/**
 * Harvest media from a single tool result.
 *
 * Uses toolParams to determine:
 * - Which extractor to use
 * - What actions to perform
 * - How to inject markdown
 *
 * @returns HarvestResult or undefined if nothing to harvest
 */
export async function harvestToolResult(
  toolContent: unknown,
  toolInfo: ToolCallInfo,
  ctx: HarvestContext,
  previewOpts: PreviewOptions
): Promise<HarvestResult | undefined> {
  const { toolName, pluginId } = toolInfo;
  const { chatWd, state, toggles, debug: ctxDebug } = ctx;

  const debug = !!ctxDebug;
  if (debug) {
    console.info(
      `[ToolResultHarvester] ENTRY: toolName="${toolName}" pluginId="${pluginId}" chatWd="${chatWd}"`
    );
  }

  // Skip if no toolName
  if (!toolName) {
    console.info(`[ToolResultHarvester] EXIT: no toolName`);
    return undefined;
  }

  // Lookup: Which mediaType handles this tool?
  // For Images (allow: "all"), we try even without pluginId using a placeholder
  const effectivePluginId = pluginId ?? "__unknown__";
  const lookup = getToolExtractor(effectivePluginId, toolName);
  if (!lookup) {
    if (debug) {
      console.info(
        `[ToolResultHarvester] No extractor for: ${effectivePluginId}/${toolName}`
      );
    }
    return undefined;
  }

  const { mediaType, extractor } = lookup;

  // Get full config for this mediaType
  const config = getMediaTypeFullConfig(mediaType);

  // Extract candidates
  // Note: The extractor itself filters out our own variant patterns (generated-image-*)
  const candidates = extractor.extractCandidates(toolContent);
  if (candidates.length === 0) {
    if (debug) {
      console.info(
        `[ToolResultHarvester] No candidates from: ${toolName} (${mediaType})`
      );
    }
    return undefined;
  }

  if (debug) {
    console.info(
      `[ToolResultHarvester] ${toolName} → ${mediaType}: ${candidates.length} candidates`
    );
  }

  // Resolve actions with toggle overrides
  const actions = resolveActions(config, toggles);

  // Resolve injection with toggle overrides
  const injection = resolveInjection(config, toggles);

  // Track state changes
  let stateChanged = false;

  // Candidates used for markdown injection (may differ from extraction candidates)
  let markdownCandidates: ToolResultCandidate[] = candidates;

  // Process based on mediaType
  if (mediaType === "picture") {
    // Pictures: External URLs, need to materialize
    const pictureRecords = await materializePictureCandidates(
      candidates,
      chatWd,
      toolInfo,
      previewOpts,
      actions.generatePreview, // Pass the flag to control preview generation
      debug
    );

    if (pictureRecords.length > 0) {
      const { changed } = appendPictures(state, pictureRecords);
      stateChanged = changed;

      const pBySourceUrl = new Map<string, number>();
      const statePictures = Array.isArray((state as any).pictures)
        ? ((state as any).pictures as any[])
        : [];
      for (const sp of statePictures) {
        const sourceUrl = typeof sp?.sourceUrl === "string" ? sp.sourceUrl : "";
        const p = typeof sp?.p === "number" ? sp.p : undefined;
        if (sourceUrl && typeof p === "number" && Number.isFinite(p)) {
          pBySourceUrl.set(sourceUrl, p);
        }
      }

      // Propagate the stable p-index back onto the original extracted candidates.
      // This enables downstream label generation (and tool-result rewrite plans)
      // to emit directly-usable notations like `pN`.
      const isDrawThingsIndex =
        toolInfo.pluginId === "ceveyne/draw-things-index" &&
        toolInfo.toolName === "index_image";

      for (const c of candidates) {
        const key = isDrawThingsIndex
          ? (typeof c.filename === "string" ? c.filename.trim() : "")
          : canonicalizeExternalUrl(
              typeof c.sourceUrl === "string"
                ? c.sourceUrl
                : typeof c.url === "string"
                  ? c.url
                  : ""
            );

        const p = key ? pBySourceUrl.get(key) : undefined;
        if (typeof p === "number" && Number.isFinite(p)) {
          c.index = p;
        }
      }

      // Prefer local previews/originals for markdown rendering so LM Studio can display them.
      // If generatePreview was false, preview/filename will be empty - renderer handles this.
      markdownCandidates = pictureRecords.map((p) => ({
        url: p.preview || undefined, // local preview in chatWd (empty if generatePreview=false)
        filename: p.filename || undefined, // local original in chatWd (empty if generatePreview=false)
        sourceUrl: p.sourceUrl, // external provenance or original path
        title: p.title,
        width: p.width,
        height: p.height,
        index: pBySourceUrl.get(p.sourceUrl),
        prompt: p.prompt,
        negativePrompt: p.negativePrompt,
        model: p.model,
        modelDisplay: p.modelDisplay,
        modelMeta: p.modelMeta,
        loras: p.loras,
        matchType: p.matchType,
        sourceInfo: p.sourceInfo,
        steps: p.steps,
        guidance: p.guidance,
        seed: p.seed,
        scheduler: p.scheduler,
        timestamp: p.timestamp,
        score: p.score,
        numFrames: p.numFrames,
        videoPath: p.videoPath || undefined,
      }));
    }
  } else if (mediaType === "image") {
    // Images: Local files already in ChatWd
    // ALWAYS generate normalized previews using central previewGenerator
    const previewSpec = config.preview;

    // Let appendImages assign indices from state.counters.nextImageI
    const imageRecords: Array<{
      filename: string;
      preview: string;
      sourceTool?: string;
      createdAt: string;
    }> = [];

    for (let i = 0; i < candidates.length; i++) {
      const c = candidates[i];
      const filename = c.filename ?? `image-${i}.png`;
      const srcAbs = path.join(chatWd, filename);
      const preview = previewFilenameFrom(filename);

      // Use central preview generator
      try {
        await generatePreview(srcAbs, chatWd, previewSpec, { debug });
      } catch (e) {
        // generatePreview logs warnings in debug mode, continue with record
      }

      // Build sourceTool in "pluginId, toolName" format
      const sourceTool =
        toolInfo.pluginId && toolInfo.toolName
          ? `${toolInfo.pluginId}, ${toolInfo.toolName}`
          : toolInfo.toolName;

      imageRecords.push({
        // Field order: filename, preview, sourceTool, (kind), (turnId), createdAt, i
        // i: omitted - appendImages assigns from state.counters.nextImageI
        // turnId: omitted - will be set by reconcile when SSOT is available
        filename,
        preview,
        sourceTool,
        createdAt: new Date().toISOString(),
      });
    }

    const { changed, records } = appendImages(state, imageRecords);
    stateChanged = changed;

    // Stabilize image indices (i-values) and inject previews (not originals)
    const recordByFilename = new Map<string, { i: number; preview: string }>();
    for (const r of records) {
      if (typeof r?.filename === "string" && typeof r?.preview === "string") {
        recordByFilename.set(r.filename, { i: r.i, preview: r.preview });
      }
    }

    for (const c of candidates) {
      const fn = typeof c.filename === "string" ? c.filename : "";
      const rec = fn ? recordByFilename.get(fn) : undefined;
      if (rec) c.index = rec.i;
    }

    markdownCandidates = records.map((r) => ({
      filename: r.preview,
      index: r.i,
    }));
  }

  // Build markdown if injection is enabled
  let markdown: string | undefined;
  if (injection.format !== "none") {
    markdown =
      extractor.renderMarkdownTable?.(markdownCandidates) ??
      renderDefaultMarkdown(markdownCandidates, injection);
  }

  return {
    mediaType,
    candidates,
    actions,
    injection,
    config,
    stateChanged,
    markdown,
  };
}

/* ═══════════════════════════════════════════════════════════════════════════
 * PICTURE MATERIALIZATION
 * ═══════════════════════════════════════════════════════════════════════════ */

interface MaterializedPicture {
  filename: string;
  preview: string;
  sourceTool?: string;
  pluginId?: string;
  sourceUrl: string;
  title?: string;
  width?: number;
  height?: number;
  pageUrl?: string;
  prompt?: string;
  negativePrompt?: string;
  model?: string;
  modelDisplay?: string;
  modelMeta?: ToolResultCandidate["modelMeta"];
  loras?: string[];
  matchType?: string;
  sourceInfo?: {
    type: "generate_image_variant" | "attachment" | "saved_image" | "draw_things_project";
    chatId?: string;
    originalName?: string;
    filePath?: string;
    projectFile?: string;
    imageType?: string;
  };
  steps?: number;
  guidance?: number;
  seed?: number;
  scheduler?: string;
  timestamp?: string;
  score?: number;
  turnId?: number;
  numFrames?: number;
  /** Absolute path to the .mov video; set iff the source file was a video. */
  videoPath?: string;
}

/**
 * Materialize picture candidates: download images and create previews.
 * If generatePreview is false, only metadata is returned (no files created).
 */
async function materializePictureCandidates(
  candidates: ToolResultCandidate[],
  chatWd: string,
  toolInfo: ToolCallInfo,
  previewOpts: PreviewOptions,
  generatePreviewFlag: boolean,
  debug?: boolean
): Promise<MaterializedPicture[]> {
  console.info(`[ToolResultHarvester] materializePictureCandidates: ${candidates.length} candidates, generatePreview=${generatePreviewFlag}`);
  
  // If generatePreview is false, return metadata-only records (no file I/O)
  if (!generatePreviewFlag) {
    console.info(`[ToolResultHarvester] generatePreview=false, returning metadata only`);
    return candidates.map((item) => ({
      filename: "", // No local file
      preview: "", // No preview
      sourceTool: toolInfo.toolName,
      pluginId: toolInfo.pluginId,
      sourceUrl: item.url || item.filename || "",
      title: item.title,
      prompt: item.prompt,
      negativePrompt: item.negativePrompt,
      model: item.model,
      modelDisplay: item.modelDisplay,
      modelMeta: item.modelMeta,
      loras: item.loras,
      matchType: item.matchType,
      sourceInfo: item.sourceInfo,
      steps: item.steps,
      guidance: item.guidance,
      seed: item.seed,
      scheduler: item.scheduler,
      timestamp: item.timestamp,
      score: item.score,
      numFrames: item.numFrames,
    }));
  }

  // Import dynamically to avoid circular deps
  const { materializeToolResultImageToFiles } = await import(
    "../core-bundle.mjs"
  );

  const results: MaterializedPicture[] = [];

  const isDrawThingsIndex =
    toolInfo.pluginId === "ceveyne/draw-things-index" &&
    toolInfo.toolName === "index_image";

  for (const item of candidates) {
    const localPath =
      typeof item.filename === "string" ? item.filename.trim() : "";
    const hasLocalPath = localPath && path.isAbsolute(localPath);
    const hasProjectUri = isProjectUri(localPath);

    const url = typeof item.url === "string" ? item.url.trim() : "";
    const externalUrl = canonicalizeExternalUrl(url);

    let sourceUrl = "";
    let originalBaseName = "";
    let previewBaseName = "";
    let originalAbs = "";
    let previewAbs = "";
    let useLocal = false;
    let useProjectUri = false;
    // Video results: set by the .mov branch below.
    let videoPath: string | undefined;       // original .mov — used as link target in the table
    let videoSourcePath: string | undefined;  // PNG last-frame sibling — copied as the "original"
    let videoPreviewPath: string | undefined; // existing preview JPG sibling — copied instead of generating

    if (isDrawThingsIndex) {
      // Handle project:// URIs (Draw Things .sqlite3 projects)
      if (hasProjectUri) {
        useProjectUri = true;
        sourceUrl = localPath; // The project:// URI
        const urlHash = shortHexSha256(localPath, 12);
        originalBaseName = `picture-h${urlHash}.png`;
        previewBaseName = `preview-picture-h${urlHash}.jpg`;
        originalAbs = path.join(chatWd, originalBaseName);
        previewAbs = path.join(chatWd, previewBaseName);
      } else if (hasLocalPath) {
        // Handle regular local files
        const localOk = await fs.promises
          .access(localPath, fs.constants.F_OK)
          .then(() => true)
          .catch(() => false);
        if (!localOk) {
          if (debug) {
            console.info(
              `[ToolResultHarvester] Skipping candidate: path not found`
            );
          }
          continue;
        }

        useLocal = true;
        sourceUrl = localPath;
        const urlHash = shortHexSha256(
          `${localPath}|${(await fs.promises.stat(localPath)).mtimeMs}`,
          12
        );
        const ext = path.extname(localPath) || ".png";
        originalBaseName = `picture-h${urlHash}${ext}`;
        previewBaseName = `preview-picture-h${urlHash}.jpg`;
        originalAbs = path.join(chatWd, originalBaseName);
        previewAbs = path.join(chatWd, previewBaseName);

        // Video files (.mov/.mp4): sharp/jimp cannot decode video, so use the PNG last-frame
        // sibling as the copyable "original" (for edits) and the existing preview JPG sibling
        // instead of attempting to generate one. sourceUrl stays as .mov for p-index tracking.
        if (/\.(mov|mp4)$/i.test(ext)) {
          const dir = path.dirname(localPath);
          const base = path.basename(localPath, ext);
          const pngSibling = path.join(dir, base + ".png");
          const previewSibling = path.join(dir, "preview-" + base + ".jpg");
          const pngOk = await fs.promises
            .access(pngSibling, fs.constants.F_OK)
            .then(() => true)
            .catch(() => false);
          if (pngOk) {
            const pngHash = shortHexSha256(
              `${pngSibling}|${(await fs.promises.stat(pngSibling)).mtimeMs}`,
              12
            );
            originalBaseName = `picture-h${pngHash}.png`;
            previewBaseName = `preview-picture-h${pngHash}.jpg`;
            originalAbs = path.join(chatWd, originalBaseName);
            previewAbs = path.join(chatWd, previewBaseName);
            videoPath = localPath; // absolute path to the original .mov — used as link target in the table
            videoSourcePath = pngSibling;
            videoPreviewPath = previewSibling;
          } else {
            // No PNG sibling — cannot render a preview for this video; skip.
            if (debug) {
              console.info(
                `[ToolResultHarvester] Skipping video: no PNG sibling for ${localPath}`
              );
            }
            continue;
          }
        }
      } else {
        if (debug) {
          console.info(
            `[ToolResultHarvester] Skipping candidate: missing path`
          );
        }
        continue;
      }
    } else {
      sourceUrl = externalUrl;
      if (!sourceUrl) {
        if (debug) {
          console.info(`[ToolResultHarvester] Skipping candidate: missing url`);
        }
        continue;
      }

      const urlHash = shortHexSha256(sourceUrl, 12);
      originalBaseName = `picture-h${urlHash}.png`;
      previewBaseName = `preview-picture-h${urlHash}.jpg`;
      originalAbs = path.join(chatWd, originalBaseName);
      previewAbs = path.join(chatWd, previewBaseName);
    }

    try {
      // Check if files already exist
      const originalOk = await fs.promises
        .access(originalAbs, fs.constants.F_OK)
        .then(() => true)
        .catch(() => false);
      const previewOk = await fs.promises
        .access(previewAbs, fs.constants.F_OK)
        .then(() => true)
        .catch(() => false);

      if (!originalOk || !previewOk) {
        if (useProjectUri) {
          // Resolve project:// URI to image buffer
          console.info(`[ToolResultHarvester] Resolving project URI: ${localPath}`);
          const imgBuffer = await resolveProjectUri(localPath);
          if (!imgBuffer) {
            console.warn(
              `[ToolResultHarvester] Failed to resolve project URI: ${localPath}`
            );
            continue;
          }
          // Write thumbnail as original (it's already JPEG/PNG)
          if (!originalOk) {
            console.info(`[ToolResultHarvester] Writing original: ${originalAbs}`);
            await fs.promises.writeFile(originalAbs, imgBuffer);
          }
          console.info(`[ToolResultHarvester] Generating preview: ${previewBaseName}`);
          await generatePreview(originalAbs, chatWd, previewOpts, {
            customFilename: previewBaseName,
            force: true,
          });
          console.info(`[ToolResultHarvester] Preview generated: ${previewBaseName}`);
        } else if (useLocal) {
          if (!originalOk) {
            await fs.promises.copyFile(videoSourcePath ?? localPath, originalAbs);
          }
          // Video: copy the .mov itself into chatWd alongside the PNG sibling.
          // Update videoPath to the local basename so MD links resolve within chatWd.
          if (videoPath) {
            const movBaseName = path.basename(videoPath);
            const movDestAbs = path.join(chatWd, movBaseName);
            const movDestOk = await fs.promises
              .access(movDestAbs, fs.constants.F_OK)
              .then(() => true)
              .catch(() => false);
            if (!movDestOk) {
              await fs.promises.copyFile(videoPath, movDestAbs);
            }
            videoPath = movBaseName;
          }
          if (videoPreviewPath) {
            // Video: copy the existing preview JPG rather than trying to encode the .mov.
            if (!previewOk) {
              const siblingOk = await fs.promises
                .access(videoPreviewPath, fs.constants.F_OK)
                .then(() => true)
                .catch(() => false);
              if (siblingOk) {
                await fs.promises.copyFile(videoPreviewPath, previewAbs);
              } else {
                // Preview sibling missing — fall back to generating from the PNG copy.
                await generatePreview(originalAbs, chatWd, previewOpts, {
                  customFilename: previewBaseName,
                  force: true,
                });
              }
            }
          } else {
            await generatePreview(originalAbs, chatWd, previewOpts, {
              customFilename: previewBaseName,
              force: true,
            });
          }
        } else {
          await materializeToolResultImageToFiles({
            url: sourceUrl,
            originalAbs,
            previewAbs,
            preview: {
              maxDim: previewOpts.maxDim,
              quality: previewOpts.quality,
            },
          });
        }
      }

      results.push({
        filename: originalBaseName,
        preview: previewBaseName,
        sourceTool: toolInfo.toolName,
        pluginId: toolInfo.pluginId,
        sourceUrl,
        title: item.title,
        width: item.width,
        height: item.height,
        pageUrl: item.sourceUrl,
        prompt: item.prompt,
        negativePrompt: item.negativePrompt,
        model: item.model,
        modelDisplay: item.modelDisplay,
        modelMeta: item.modelMeta,
        loras: item.loras,
        matchType: item.matchType,
        sourceInfo: item.sourceInfo,
        steps: item.steps,
        guidance: item.guidance,
        seed: item.seed,
        scheduler: item.scheduler,
        timestamp: item.timestamp,
        score: item.score,
        numFrames: item.numFrames,
        videoPath,
      });
    } catch (e) {
      console.warn(
        `[ToolResultHarvester] Failed to materialize ${sourceUrl}:`,
        (e as Error).message
      );
    }
  }

  return results;
}

/* ═══════════════════════════════════════════════════════════════════════════
 * MARKDOWN RENDERING
 * ═══════════════════════════════════════════════════════════════════════════ */

/**
 * Default markdown renderer based on injection config.
 */
function renderDefaultMarkdown(
  candidates: ToolResultCandidate[],
  injectMdInAgentResponse: InjectionConfig
): string {
  if (injectMdInAgentResponse.format === "none") {
    return "";
  }

  if (injectMdInAgentResponse.format === "markdown-table") {
    // Table format with header
    const header =
      injectMdInAgentResponse.tableHeader ?? "| Preview | Title | Source |";
    const separator = "|" + "-|".repeat((header.match(/\|/g)?.length ?? 3) - 1);

    const rows = candidates.map((item, i) => {
      const template =
        injectMdInAgentResponse.itemTemplate ??
        "| ![]({url}) | {title} | {source} |";
      return template
        .replace("{url}", item.url ?? "")
        .replace("{title}", item.title ?? "")
        .replace("{source}", item.sourceUrl ?? item.url ?? "")
        .replace("{preview}", item.filename ?? "")
        .replace("{index}", String(i + 1));
    });

    return [header, separator, ...rows].join("\n");
  }

  if (injectMdInAgentResponse.format === "markdown-image") {
    // Simple image list
    return candidates
      .map((item, i) => {
        const label = injectMdInAgentResponse.labelGenerator(item, i);
        const path = item.filename ?? item.url ?? "";
        return `![${label}](./${path})`;
      })
      .join("\n");
  }

  if (injectMdInAgentResponse.format === "markdown-link") {
    // Simple link list
    return candidates
      .map((item, i) => {
        const label = injectMdInAgentResponse.labelGenerator(item, i);
        const path = item.filename ?? item.url ?? "";
        return `[${label}](./${path})`;
      })
      .join("\n");
  }

  return "";
}

/* ═══════════════════════════════════════════════════════════════════════════
 * BATCH HARVESTING
 * ═══════════════════════════════════════════════════════════════════════════ */

/**
 * Harvest all tool results from a set of tool calls.
 * Returns combined results for injection.
 */
export async function harvestAllToolResults(
  toolResults: Array<{ content: unknown; info: ToolCallInfo }>,
  ctx: HarvestContext,
  previewOpts: PreviewOptions
): Promise<HarvestResult[]> {
  const results: HarvestResult[] = [];

  for (const { content, info } of toolResults) {
    const result = await harvestToolResult(content, info, ctx, previewOpts);
    if (result) {
      results.push(result);
    }
  }

  return results;
}

/**
 * Inject markdown for all harvest results.
 * @deprecated This function is not currently used. Injection tracking is handled
 * at the orchestrator level with trackInjection() which requires chatId.
 */
export function injectHarvestMarkdown(
  results: HarvestResult[],
  ctl: GeneratorController
): void {
  for (const result of results) {
    if (
      result.markdown &&
      result.config.injectMdInAgentResponse.format !== "none"
    ) {
      // Note: No injection tracking here - handled by orchestrator
      ctl.fragmentGenerated(result.markdown);
    }
  }
}