analyse-image

import fs from "fs";

/* -------------------------------------------------------------------------- */
/* PNG XMP metadata reader                                                     */
/* Reads the exif:UserComment JSON from a Draw Things iTXt/XMP chunk.         */
/* -------------------------------------------------------------------------- */

export interface PngGenerationMeta {
  prompt?: string;
  negativePrompt?: string;
  model?: string;
  sampler?: string;
  steps?: number;
  guidanceScale?: number;
  seed?: number;
  seedMode?: string;
  shift?: number;
  size?: string;
  strength?: number;
  loras?: Array<{ file: string; weight?: number }>;
  sources?: string[];
  mode?: string;
  generatedBy?: string;
}

/**
 * Read Draw Things generation metadata from a PNG's XMP iTXt chunk.
 * Returns null if the file is not a PNG, has no XMP chunk, or the JSON cannot
 * be parsed.
 */
export function readPngGenerationMeta(filePath: string): PngGenerationMeta | null {
  let buf: Buffer;
  try {
    buf = fs.readFileSync(filePath);
  } catch {
    return null;
  }

  // Verify PNG signature: \x89 P N G \r \n \x1a \n
  if (buf.length < 8 || buf[0] !== 0x89 || buf[1] !== 0x50 || buf[2] !== 0x4e || buf[3] !== 0x47) {
    return null;
  }

  let offset = 8; // skip 8-byte PNG signature
  while (offset + 12 <= buf.length) {
    const chunkLen = buf.readUInt32BE(offset);
    const chunkType = buf.toString("ascii", offset + 4, offset + 8);
    const dataStart = offset + 8;
    const dataEnd = dataStart + chunkLen;

    if (dataEnd + 4 > buf.length) break;

    if (chunkType === "IEND") break;

    if (chunkType === "iTXt") {
      const data = buf.slice(dataStart, dataEnd);
      const kwEnd = data.indexOf(0);
      if (kwEnd >= 0 && data.toString("ascii", 0, kwEnd) === "XML:com.adobe.xmp") {
        const comprFlag = data[kwEnd + 1]; // 0 = uncompressed, 1 = compressed
        if (comprFlag !== 0) {
          offset = dataEnd + 4;
          continue;
        }
        // Skip: comprFlag(1) + comprMethod(1) = offset + 2 after kwEnd
        let pos = kwEnd + 3;
        // Skip language tag (null-terminated)
        while (pos < data.length && data[pos] !== 0) pos++;
        pos++;
        // Skip translated keyword (null-terminated)
        while (pos < data.length && data[pos] !== 0) pos++;
        pos++;
        // Remaining bytes are the XMP text
        const xmpText = data.toString("utf8", pos);
        return extractMetaFromXmp(xmpText);
      }
    }

    offset = dataEnd + 4; // advance past chunk + 4-byte CRC
  }

  return null;
}

function extractMetaFromXmp(xmp: string): PngGenerationMeta | null {
  // The JSON lives inside <exif:UserComment><rdf:Alt><rdf:li ...>JSON</rdf:li></rdf:Alt></exif:UserComment>
  // (Draw Things does NOT use CDATA here.)
  const match = xmp.match(/<exif:UserComment>[\s\S]*?<rdf:li[^>]*>([\s\S]*?)<\/rdf:li>/);
  if (!match) return null;

  let raw: Record<string, any>;
  try {
    // XML text content may have been escaped (e.g. &quot; for ") by older writer builds.
    // Unescape the five standard XML entities before JSON.parse.
    const jsonText = match[1].trim()
      .replace(/&quot;/g, '"')
      .replace(/&apos;/g, "'")
      .replace(/&lt;/g, "<")
      .replace(/&gt;/g, ">")
      .replace(/&amp;/g, "&");
    raw = JSON.parse(jsonText);
  } catch {
    return null;
  }

  // Draw Things writes key "model"; our pngMetadata.ts writer also uses "model".
  const loras: PngGenerationMeta["loras"] = Array.isArray(raw.lora)
    ? raw.lora
        .filter((l: any) => l && typeof l.model === "string")
        .map((l: any) => ({
          file: String(l.model),
          weight: typeof l.weight === "number" ? l.weight : undefined,
        }))
    : undefined;

  const sources: string[] | undefined = Array.isArray(raw.sources)
    ? raw.sources.filter((s: any) => typeof s === "string")
    : undefined;

  return {
    prompt: typeof raw.c === "string" && raw.c ? raw.c : undefined,
    negativePrompt: typeof raw.uc === "string" && raw.uc ? raw.uc : undefined,
    model: typeof raw.model === "string" && raw.model ? raw.model : undefined,
    sampler: typeof raw.sampler === "string" && raw.sampler ? raw.sampler : undefined,
    steps: typeof raw.steps === "number" ? raw.steps : undefined,
    guidanceScale: typeof raw.scale === "number" ? raw.scale : undefined,
    seed: typeof raw.seed === "number" ? raw.seed : undefined,
    seedMode: typeof raw.seed_mode === "string" && raw.seed_mode ? raw.seed_mode : undefined,
    shift: typeof raw.shift === "number" ? raw.shift : undefined,
    size: typeof raw.size === "string" && raw.size ? raw.size : undefined,
    strength: typeof raw.strength === "number" ? raw.strength : undefined,
    loras: loras?.length ? loras : undefined,
    sources: sources?.length ? sources : undefined,
    mode: typeof raw.mode === "string" && raw.mode ? raw.mode : undefined,
    generatedBy: typeof raw.generated_by === "string" && raw.generated_by ? raw.generated_by : undefined,
  };
}

export function formatGenerationMeta(meta: PngGenerationMeta): string {
  const lines: string[] = ["  GENERATION METADATA:"];
  if (meta.prompt) lines.push(`  Prompt: ${meta.prompt}`);
  if (meta.negativePrompt) lines.push(`  Negative Prompt: ${meta.negativePrompt}`);
  if (meta.model) lines.push(`  Model: ${meta.model}`);

  const techParts: string[] = [];
  if (meta.sampler) techParts.push(`Sampler: ${meta.sampler}`);
  if (typeof meta.steps === "number") techParts.push(`Steps: ${meta.steps}`);
  if (typeof meta.guidanceScale === "number") techParts.push(`Guidance Scale: ${meta.guidanceScale}`);
  if (typeof meta.seed === "number") techParts.push(`Seed: ${meta.seed}`);
  if (techParts.length > 0) lines.push(`  ${techParts.join("  ")}`);

  if (meta.size) lines.push(`  Size: ${meta.size}`);
  if (typeof meta.strength === "number" && meta.strength !== 1.0) {
    lines.push(`  Strength: ${meta.strength}`);
  }
  if (meta.loras && meta.loras.length > 0) {
    const loraStr = meta.loras
      .map((l) => (l.weight != null ? `${l.file} (${l.weight})` : l.file))
      .join(", ");
    lines.push(`  LoRA: ${loraStr}`);
  }
  if (meta.sources && meta.sources.length > 0) {
    lines.push(`  Source(s): ${meta.sources.join(", ")}`);
  }
  return lines.join("\n");
}