draw-things-index

import fs from "node:fs";
import path from "node:path";

export type MediaKind = "attachment" | "variant" | "image" | "picture";

export interface ResolvedMediaQuery {
  notation: string;
  mediaKind: MediaKind;
  originalName?: string;
  filePath?: string;
  previewPath?: string;
  displayName: string;
}

export interface MediaResolutionResult {
  rewrittenQuery: string;
  mediaQueries: ResolvedMediaQuery[];
  notices: string[];
}

const IMAGE_FILENAME_EXTS = new Set(["png", "jpg", "jpeg", "webp", "gif", "tif", "tiff", "bmp", "heic", "heif"]);

function isImagePath(value: string): boolean {
  const ext = path.extname(value).slice(1).toLowerCase();
  return IMAGE_FILENAME_EXTS.has(ext);
}

function absoluteIfExists(workingDir: string, value: unknown): string | undefined {
  if (typeof value !== "string" || !value.trim()) return undefined;
  const expanded = value.startsWith("~") ? path.join(process.env.HOME ?? "", value.slice(1)) : value;
  const candidate = path.isAbsolute(expanded) ? expanded : path.join(workingDir, expanded);
  try {
    if (fs.existsSync(candidate) && fs.statSync(candidate).isFile() && isImagePath(candidate)) {
      return candidate;
    }
  } catch {
    return undefined;
  }
  return undefined;
}

function displayNameFor(record: any, candidates: Array<unknown>): string | undefined {
  if (typeof record?.originalName === "string" && record.originalName.trim()) return record.originalName;
  for (const candidate of candidates) {
    if (typeof candidate !== "string" || !candidate.trim()) continue;
    try {
      if (/^https?:\/\//i.test(candidate)) return candidate;
      return path.basename(candidate);
    } catch {
      return candidate;
    }
  }
  return undefined;
}

function resolveAttachment(workingDir: string, record: any): ResolvedMediaQuery | null {
  if (typeof record?.a !== "number" || record.a <= 0) return null;
  const originalPath = absoluteIfExists(workingDir, record.originAbs) ?? absoluteIfExists(workingDir, record.filename);
  const previewPath = absoluteIfExists(workingDir, record.preview);
  const displayName = displayNameFor(record, [record.filename, record.originAbs, record.preview]);
  if (!displayName) return null;
  return {
    notation: `a${record.a}`,
    mediaKind: "attachment",
    originalName: typeof record.originalName === "string" ? record.originalName : undefined,
    filePath: originalPath ?? previewPath,
    previewPath,
    displayName,
  };
}

function resolveGenerated(workingDir: string, record: any, key: "v" | "i", mediaKind: "variant" | "image"): ResolvedMediaQuery | null {
  if (typeof record?.[key] !== "number" || record[key] <= 0) return null;
  const originalPath = absoluteIfExists(workingDir, record.filename);
  const previewPath = absoluteIfExists(workingDir, record.preview);
  const displayName = displayNameFor(record, [record.filename, record.preview, record.sourceUrl]);
  if (!displayName) return null;
  return {
    notation: `${key}${record[key]}`,
    mediaKind,
    filePath: originalPath ?? previewPath,
    previewPath,
    displayName,
  };
}

function resolvePicture(workingDir: string, record: any): ResolvedMediaQuery | null {
  if (typeof record?.p !== "number" || record.p <= 0) return null;
  const originalPath = absoluteIfExists(workingDir, record.filename);
  const previewPath = absoluteIfExists(workingDir, record.preview);
  const displayName = displayNameFor(record, [record.filename, record.preview, record.sourceUrl]);
  if (!displayName) return null;
  return {
    notation: `p${record.p}`,
    mediaKind: "picture",
    filePath: originalPath ?? previewPath,
    previewPath,
    displayName,
  };
}

export async function resolveMediaQueries(workingDir: string, query: string): Promise<MediaResolutionResult> {
  const stateFile = path.join(workingDir, "chat_media_state.json");
  const raw = await fs.promises.readFile(stateFile, "utf8");
  const state: any = JSON.parse(raw);
  const index = new Map<string, ResolvedMediaQuery>();

  for (const record of Array.isArray(state?.attachments) ? state.attachments : []) {
    const resolved = resolveAttachment(workingDir, record);
    if (resolved) index.set(resolved.notation.toLowerCase(), resolved);
  }
  for (const record of Array.isArray(state?.variants) ? state.variants : []) {
    const resolved = resolveGenerated(workingDir, record, "v", "variant");
    if (resolved) index.set(resolved.notation.toLowerCase(), resolved);
  }
  for (const record of Array.isArray(state?.images) ? state.images : []) {
    const resolved = resolveGenerated(workingDir, record, "i", "image");
    if (resolved) index.set(resolved.notation.toLowerCase(), resolved);
  }
  for (const record of Array.isArray(state?.pictures) ? state.pictures : []) {
    const resolved = resolvePicture(workingDir, record);
    if (resolved) index.set(resolved.notation.toLowerCase(), resolved);
  }

  const mediaQueries: ResolvedMediaQuery[] = [];
  const notices: string[] = [];
  const rewrittenQuery = query.replace(/\b([avip]\d+)\b/gi, (match) => {
    const resolved = index.get(match.toLowerCase());
    if (!resolved) {
      notices.push(`Media notation ${match} was not found in chat_media_state.json.`);
      return match;
    }
    mediaQueries.push(resolved);
    if (!resolved.filePath) {
      notices.push(`Media notation ${match} resolved to ${resolved.displayName}, but no local image file is available.`);
    }
    return resolved.displayName;
  });

  return { rewrittenQuery, mediaQueries, notices };
}