user-docs

import fs from "node:fs";
import path from "node:path";
import { createHash } from "node:crypto";
import {
  readState,
  writeStateAtomic,
  generatePreviewFromBuffer,
  materializeToolResultImageToFiles,
  appendPictures,
  getSelfPluginIdentifier,
} from "../core-bundle.mjs";
import { parseRemoteImageRefsFromChunk } from "../sources/remoteImageResolver.js";

export interface RegisteredDocumentImages {
  count: number;
  assignedKeys: string[];
  assignedDisplayLines: string[];
  remoteSkips: Record<string, number>;
}

export interface RegisterDocumentImagesOptions {
  markdown: string;
  documentPath: string;
  chatWd: string;
  sourceTool: string;
  sourceKind?: string;
  chunkMetadata?: Record<string, unknown>;
  maxImages?: number;
  previewMaxSum: number;
  previewQuality: number;
  remoteFetchTimeoutMs?: number;
  remoteMaxBytes?: number;
  onStatus?: (message: string) => void;
}

export interface RegisterDocumentImageFilesOptions {
  imagePaths: string[];
  chatWd: string;
  sourceTool: string;
  previewMaxSum: number;
  previewQuality: number;
  onStatus?: (message: string) => void;
}

export function sanitizeEmbeddedImagePayloads(input: string, replacement = "[embedded image -- registered]"): string {
  let output = String(input ?? "");
  output = output.replace(
    /data:image\/[a-zA-Z0-9+.-]+(?:;[a-zA-Z0-9_.+-]+=[^;,\s)]*)*;base64,[A-Za-z0-9+/=\r\n]+/g,
    replacement
  );
  output = output.replace(/^\s*[A-Za-z0-9+/=]{256,}\s*$/gm, replacement);
  return output;
}

export function parseMarkdownImageRefs(
  chunkContent: string,
  documentPath: string,
  tempDir?: string
): string[] {
  const docDir = path.dirname(documentPath);
  const refs: string[] = [];
  const seen = new Set<string>();

  function addAbs(abs: string | null): void {
    if (abs && !seen.has(abs)) {
      seen.add(abs);
      refs.push(abs);
    }
  }

  const reAngle = /!\[.*?\]\(<([^>]+)>\)/g;
  const rePlain = /!\[.*?\]\((?!<)(?!data:)([^)\s"]+)(?:\s+"[^"]*")?\)/g;

  for (const re of [reAngle, rePlain]) {
    let match: RegExpExecArray | null;
    while ((match = re.exec(chunkContent)) !== null) {
      const raw = match[1].trim();
      if (/^https?:\/\//i.test(raw) || /^data:/i.test(raw)) continue;
      const abs = /^file:\/\//i.test(raw) ? fileUrlToPath(raw) : path.resolve(docDir, raw);
      if (abs && fs.existsSync(abs)) addAbs(abs);
    }
  }

  const reWiki = /!\[\[([^\]]+)\]\]/g;
  let wikiMatch: RegExpExecArray | null;
  while ((wikiMatch = reWiki.exec(chunkContent)) !== null) {
    const raw = wikiMatch[1].split("|")[0].trim();
    if (/^https?:\/\//i.test(raw)) continue;
    const abs = /^file:\/\//i.test(raw) ? fileUrlToPath(raw) : path.resolve(docDir, raw);
    if (abs && fs.existsSync(abs)) addAbs(abs);
  }

  if (!tempDir) return refs;

  const reB64Md = /!\[.*?\]\((?:<)?data:image\/([a-zA-Z0-9+.-]+);base64,([^)>]+)(?:>)?\)/g;
  let b64Match: RegExpExecArray | null;
  while ((b64Match = reB64Md.exec(chunkContent)) !== null) {
    addAbs(decodeBase64ToFile(b64Match[2].replace(/\s/g, ""), b64Match[1].toLowerCase(), documentPath, tempDir));
  }

  const reB64Html = /<img\s[^>]*src=["']data:image\/([a-zA-Z0-9+.-]+);base64,([^"']+)["'][^>]*>/gi;
  let htmlMatch: RegExpExecArray | null;
  while ((htmlMatch = reB64Html.exec(chunkContent)) !== null) {
    addAbs(decodeBase64ToFile(htmlMatch[2].replace(/\s/g, ""), htmlMatch[1].toLowerCase(), documentPath, tempDir));
  }

  return refs;
}

export function extractBase64ImagesFromDocument(docPath: string, tempDir: string, content?: string): string[] {
  try {
    const source = content ?? fs.readFileSync(docPath, "utf8");
    const refs: string[] = [];
    const seen = new Set<string>();
    function add(abs: string | null): void {
      if (abs && !seen.has(abs)) {
        seen.add(abs);
        refs.push(abs);
      }
    }

    const reB64Md = /!\[.*?\]\(<?data:image\/([a-zA-Z0-9+.-]+);base64,([^>)]+)>?\)/g;
    let match: RegExpExecArray | null;
    while ((match = reB64Md.exec(source)) !== null) {
      add(decodeBase64ToFile(match[2].replace(/\s/g, ""), match[1].toLowerCase(), docPath, tempDir));
    }

    const reHtml = /<img\s[^>]*src=["']data:image\/([a-zA-Z0-9+.-]+);base64,([^"']+)["'][^>]*>/gi;
    while ((match = reHtml.exec(source)) !== null) {
      add(decodeBase64ToFile(match[2].replace(/\s/g, ""), match[1].toLowerCase(), docPath, tempDir));
    }

    return refs;
  } catch {
    return [];
  }
}

export function localImageRefsFromMetadata(value: unknown): string[] {
  if (!Array.isArray(value)) return [];
  const refs: string[] = [];
  const seen = new Set<string>();
  for (const item of value) {
    const raw = typeof item === "string"
      ? item
      : item && typeof item === "object" && typeof (item as any).url === "string"
        ? (item as any).url
        : "";
    const abs = localImageRefToPath(raw);
    if (!abs || seen.has(abs) || !fs.existsSync(abs)) continue;
    seen.add(abs);
    refs.push(abs);
  }
  return refs;
}

export async function registerDocumentImagesFromMarkdown(
  options: RegisterDocumentImagesOptions
): Promise<RegisteredDocumentImages> {
  const maxImages = options.maxImages ?? Number.MAX_SAFE_INTEGER;
  const imagePathSet = new Set<string>();
  const remoteImageRefs: Array<{ url: string; altText?: string }> = [];
  const remoteImageUrlSet = new Set<string>();

  const sourceKind = options.sourceKind ?? "file";
  if (sourceKind === "github" || sourceKind === "huggingface" || sourceKind === "https") {
    const refs = parseRemoteImageRefsFromChunk(options.markdown, options.chunkMetadata ?? {});
    for (const ref of refs) {
      if (!remoteImageUrlSet.has(ref.url)) {
        remoteImageUrlSet.add(ref.url);
        remoteImageRefs.push({ url: ref.url, altText: ref.altText });
      }
      if (remoteImageRefs.length >= maxImages) break;
    }
  } else {
    for (const abs of parseMarkdownImageRefs(options.markdown, options.documentPath, options.chatWd)) {
      imagePathSet.add(abs);
      if (imagePathSet.size >= maxImages) break;
    }

    if (options.markdown.includes("data:image/") && imagePathSet.size < maxImages) {
      for (const abs of extractBase64ImagesFromDocument(options.documentPath, options.chatWd, options.markdown)) {
        imagePathSet.add(abs);
        if (imagePathSet.size >= maxImages) break;
      }
    }

    for (const abs of localImageRefsFromMetadata((options.chunkMetadata as any)?.imageRefs)) {
      imagePathSet.add(abs);
      if (imagePathSet.size >= maxImages) break;
    }
  }

  const imagePaths = Array.from(imagePathSet).slice(0, maxImages);
  if (imagePaths.length === 0 && remoteImageRefs.length === 0) {
    return { count: 0, assignedKeys: [], assignedDisplayLines: [], remoteSkips: {} };
  }

  await fs.promises.mkdir(options.chatWd, { recursive: true }).catch(() => {});

  const localRecords = await buildPictureRecordsFromLocalFiles({
    imagePaths,
    chatWd: options.chatWd,
    sourceTool: options.sourceTool,
    previewMaxSum: options.previewMaxSum,
    previewQuality: options.previewQuality,
    onStatus: options.onStatus,
  });

  const pluginId = getSelfPluginIdentifier() ?? "unknown";
  const imageRecords: any[] = [...localRecords];
  const remoteSkips: Record<string, number> = {};

  for (let index = 0; index < remoteImageRefs.length && imageRecords.length < maxImages; index++) {
    const ref = remoteImageRefs[index];
    try {
      options.onStatus?.(`Fetching remote image ${index + 1}/${remoteImageRefs.length}...`);
      const originalBaseName = sanitizeImageFilename(filenameFromRemoteUrl(ref.url));
      const previewBaseName = previewFilenameFor(originalBaseName);
      const originalAbs = path.join(options.chatWd, originalBaseName);
      const previewAbs = path.join(options.chatWd, previewBaseName);
      await materializeToolResultImageToFiles({
        url: ref.url,
        originalAbs,
        previewAbs,
        preview: {
          maxDim: options.previewMaxSum,
          quality: options.previewQuality,
        },
        timeoutMs: options.remoteFetchTimeoutMs ?? 10000,
        maxBytes: options.remoteMaxBytes ?? 15728640,
      });
      imageRecords.push({
        filename: originalBaseName,
        preview: previewBaseName,
        sourceTool: `${pluginId}/${options.sourceTool}`,
        sourceUrl: ref.url,
        title: ref.altText,
      });
    } catch (err) {
      const reason = err instanceof Error ? err.message : String(err);
      remoteSkips[reason] = (remoteSkips[reason] ?? 0) + 1;
      console.warn(`[documentImages] remote image materialization failed for ${ref.url}:`, reason);
    }
  }

  return appendPictureRecords(options.chatWd, imageRecords, remoteSkips);
}

export async function registerDocumentImageFiles(
  options: RegisterDocumentImageFilesOptions
): Promise<RegisteredDocumentImages> {
  await fs.promises.mkdir(options.chatWd, { recursive: true }).catch(() => {});
  const imageRecords = await buildPictureRecordsFromLocalFiles(options);
  return appendPictureRecords(options.chatWd, imageRecords, {});
}

async function buildPictureRecordsFromLocalFiles(options: RegisterDocumentImageFilesOptions): Promise<any[]> {
  const pluginId = getSelfPluginIdentifier() ?? "unknown";
  const previewSpec = {
    maxDim: options.previewMaxSum,
    maxSum: options.previewMaxSum,
    mode: "sum" as const,
    quality: options.previewQuality,
  };
  const records: any[] = [];
  const seen = new Set<string>();

  for (let index = 0; index < options.imagePaths.length; index++) {
    const srcAbs = options.imagePaths[index];
    try {
      if (!fs.existsSync(srcAbs)) continue;
      const realSource = fs.realpathSync(srcAbs);
      if (seen.has(realSource)) continue;
      seen.add(realSource);

      options.onStatus?.(`Generating image preview ${index + 1}/${options.imagePaths.length}...`);
      const srcBuf = await fs.promises.readFile(srcAbs);
      const originalFilename = sanitizeImageFilename(path.basename(srcAbs));
      const destAbs = await copyOriginalIntoChatWd(srcAbs, options.chatWd, originalFilename, srcBuf);
      const destFilename = path.basename(destAbs);
      const preview = await generatePreviewFromBuffer(
        srcBuf,
        options.chatWd,
        destFilename,
        previewSpec,
        { customFilename: previewFilenameFor(destFilename) }
      );
      records.push({
        filename: destFilename,
        preview: preview.previewFilename,
        sourceTool: `${pluginId}/${options.sourceTool}`,
        sourceUrl: `file://${realSource}`,
      });
    } catch (err) {
      console.warn(`[documentImages] preview generation failed for ${srcAbs}:`, String(err));
    }
  }

  return records;
}

async function appendPictureRecords(
  chatWd: string,
  imageRecords: any[],
  remoteSkips: Record<string, number>
): Promise<RegisteredDocumentImages> {
  if (imageRecords.length === 0) {
    return { count: 0, assignedKeys: [], assignedDisplayLines: [], remoteSkips };
  }

  try {
    const state = await readState(chatWd);
    const appendResult = appendPictures(state, imageRecords);
    if (appendResult.changed) {
      await writeStateAtomic(chatWd, state);
    }
    const pictures: any[] = Array.isArray((state as any)?.pictures) ? (state as any).pictures : [];
    const recordsBySourceUrl = new Map<string, any>();
    for (const rec of pictures) {
      const sourceUrl = typeof rec?.sourceUrl === "string" ? rec.sourceUrl.trim() : "";
      if (sourceUrl && typeof rec?.p === "number" && !recordsBySourceUrl.has(sourceUrl)) {
        recordsBySourceUrl.set(sourceUrl, rec);
      }
    }
    const resolvedRecords = imageRecords
      .map((rec) => {
        const sourceUrl = typeof rec?.sourceUrl === "string" ? rec.sourceUrl.trim() : "";
        return sourceUrl ? recordsBySourceUrl.get(sourceUrl) : undefined;
      })
      .filter((record): record is any => !!record);
    return {
      count: imageRecords.length,
      assignedKeys: resolvedRecords
        .map((record) => (typeof record.p === "number" ? `p${record.p}` : null))
        .filter((key): key is string => key !== null),
      assignedDisplayLines: resolvedRecords
        .filter((record) => typeof record.p === "number" && typeof record.preview === "string")
        .map((record) => `p${record.p}: ![p${record.p}](./${record.preview})`),
      remoteSkips,
    };
  } catch (err) {
    console.warn("[documentImages] state update failed:", String(err));
    return { count: imageRecords.length, assignedKeys: [], assignedDisplayLines: [], remoteSkips };
  }
}

async function copyOriginalIntoChatWd(srcAbs: string, chatWd: string, preferredFilename: string, srcBuf: Buffer): Promise<string> {
  const srcResolved = path.resolve(srcAbs);
  let destAbs = path.join(chatWd, preferredFilename);
  if (path.resolve(destAbs) === srcResolved) return destAbs;
  destAbs = await chooseDeterministicDestination(destAbs, srcBuf);
  if (!fs.existsSync(destAbs)) {
    await fs.promises.writeFile(destAbs, srcBuf);
  }
  return destAbs;
}

async function chooseDeterministicDestination(preferredAbs: string, srcBuf: Buffer): Promise<string> {
  if (!fs.existsSync(preferredAbs)) return preferredAbs;
  try {
    const existing = await fs.promises.readFile(preferredAbs);
    if (existing.equals(srcBuf)) return preferredAbs;
  } catch {}
  const parsed = path.parse(preferredAbs);
  for (let index = 2; index < 1000; index++) {
    const candidate = path.join(parsed.dir, `${parsed.name}-${index}${parsed.ext}`);
    if (!fs.existsSync(candidate)) return candidate;
    try {
      const existing = await fs.promises.readFile(candidate);
      if (existing.equals(srcBuf)) return candidate;
    } catch {}
  }
  return preferredAbs;
}

function decodeBase64ToFile(b64: string, mime: string, documentPath: string, tempDir: string): string | null {
  try {
    const ext = mime === "jpeg" ? "jpg" : mime.replace(/[^a-z0-9]/gi, "");
    const stem = path.basename(documentPath, path.extname(documentPath)).replace(/[^a-z0-9_-]/gi, "_");
    const hash = createHash("sha256").update(b64).digest("hex").slice(0, 16);
    const filename = sanitizeImageFilename(`embedded-${stem}-${hash}.${ext}`);
    const abs = path.join(tempDir, filename);
    if (!fs.existsSync(abs)) {
      fs.writeFileSync(abs, Buffer.from(b64, "base64"));
    }
    return abs;
  } catch {
    return null;
  }
}

function localImageRefToPath(value: string): string | null {
  const ref = String(value ?? "").trim();
  if (!ref || /^data:/i.test(ref) || /^https?:\/\//i.test(ref)) return null;
  if (/^file:\/\//i.test(ref)) return fileUrlToPath(ref);
  return path.isAbsolute(ref) ? ref : null;
}

function sanitizeImageFilename(value: string): string {
  const parsed = path.parse(path.basename(value).replace(/ /g, "_"));
  const name = parsed.name.replace(/[\\/:*?"<>|\x00-\x1F]/g, "_") || "image";
  const ext = parsed.ext.replace(/[\\/:*?"<>|\x00-\x1F]/g, "_") || ".png";
  return `${name}${ext}`;
}

function previewFilenameFor(originalFilename: string): string {
  const parsed = path.parse(sanitizeImageFilename(originalFilename));
  return `preview-${parsed.name}.jpg`;
}

function filenameFromRemoteUrl(value: string): string {
  try {
    const url = new URL(value);
    const base = path.basename(decodeURIComponent(url.pathname));
    if (base && base !== "/" && path.extname(base)) return base;
  } catch {}
  return "remote_image.png";
}

function fileUrlToPath(value: string): string | null {
  try {
    return new URL(value).pathname;
  } catch {
    return null;
  }
}