// PNG iTXt chunk injection for Draw Things-compatible XMP metadata.
import os from "node:os";

export type PngAnalysisMetadata = {
  schema: "ceveyne.image-analysis/v1";
  tool: "detect_object" | "annotate_image";
  sourceNotation: string;
  inference?: {
    model: string;
    query?: string;
    detectorPromptSha256?: string;
    maxTokens?: number;
    temperature?: number;
    reused?: boolean;
  };
  render?: {
    palette?: boolean;
    color?: string;
    lineWeight?: number;
    frameAdjust?: number | string;
  };
  detections: Array<{
    label: string;
    bbox: { x1: number; y1: number; x2: number; y2: number };
  }>;
};

export type PngXmpParams = {
  prompt?: string;
  model?: string;
  width?: number;
  height?: number;
  sources?: string[];
  mode?: string;
  generatedBy?: string;
  creatorTool?: string;
  createdAt?: string;
  analysis?: PngAnalysisMetadata;
};

const CRC_TABLE = (() => {
  const table = new Int32Array(256);
  for (let n = 0; n < 256; n++) {
    let crc = n;
    for (let bit = 0; bit < 8; bit++) {
      crc = crc & 1 ? (0xedb88320 ^ (crc >>> 1)) : (crc >>> 1);
    }
    table[n] = crc;
  }
  return table;
})();

function crc32(data: Buffer): number {
  let crc = 0xffffffff;
  for (let index = 0; index < data.length; index++) {
    crc = CRC_TABLE[(crc ^ data[index]!) & 0xff]! ^ (crc >>> 8);
  }
  return (crc ^ 0xffffffff) >>> 0;
}

function normPath(absolutePath: string): string {
  const home = os.homedir();
  return absolutePath.startsWith(home) ? `~${absolutePath.slice(home.length)}` : absolutePath;
}

function escapeXml(value: string): string {
  return value
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;");
}

function getPngDimensions(pngBuffer: Buffer): { width: number; height: number } | undefined {
  if (
    pngBuffer.length < 24 ||
    !pngBuffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a])) ||
    pngBuffer.toString("ascii", 12, 16) !== "IHDR"
  ) {
    return undefined;
  }
  return { width: pngBuffer.readUInt32BE(16), height: pngBuffer.readUInt32BE(20) };
}

function buildXmpString(params: PngXmpParams): string {
  const createdAt = escapeXml(params.createdAt ?? new Date().toISOString());
  const metadata: Record<string, unknown> = {};
  if (params.prompt != null) metadata["c"] = params.prompt;
  if (params.model) metadata["model"] = params.model;
  if (typeof params.width === "number" && typeof params.height === "number") {
    metadata["size"] = `${params.width}x${params.height}`;
  }
  if (params.sources?.length) metadata["sources"] = params.sources.map(normPath);
  if (params.mode) metadata["mode"] = params.mode;
  if (params.analysis) metadata["analysis"] = params.analysis;
  metadata["generated_by"] = params.generatedBy ?? "ceveyne/analyse-image";

  const description: string[] = [];
  if (params.prompt) description.push(params.prompt);
  const details: string[] = [];
  if (typeof params.width === "number" && typeof params.height === "number") details.push(`Size: ${params.width}x${params.height}`);
  if (params.model) details.push(`Model: ${params.model}`);
  if (details.length) description.push(details.join(", "));
  if (params.sources?.length) description.push(`Source: ${params.sources.map(normPath).join(", ")}`);

  return [
    `<x:xmpmeta xmlns:x="adobe:ns:meta/" x:xmptk="XMP Core 6.0.0">`,
    `   <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">`,
    `      <rdf:Description rdf:about=""`,
    `            xmlns:dc="http://purl.org/dc/elements/1.1/"`,
    `            xmlns:xmp="http://ns.adobe.com/xap/1.0/"`,
    `            xmlns:exif="http://ns.adobe.com/exif/1.0/">`,
    `         <dc:description><rdf:Alt><rdf:li xml:lang="x-default">${escapeXml(description.join("\n"))}</rdf:li></rdf:Alt></dc:description>`,
    `         <xmp:CreatorTool>${escapeXml(params.creatorTool ?? params.generatedBy ?? "ceveyne/analyse-image")}</xmp:CreatorTool>`,
    `         <xmp:CreateDate>${createdAt}</xmp:CreateDate>`,
    `         <exif:UserComment><rdf:Alt><rdf:li xml:lang="x-default">${JSON.stringify(metadata)}</rdf:li></rdf:Alt></exif:UserComment>`,
    `      </rdf:Description>`,
    `   </rdf:RDF>`,
    `</x:xmpmeta>`,
  ].join("\n");
}

function buildITxtChunk(text: string): Buffer {
  const data = Buffer.concat([
    Buffer.from("XML:com.adobe.xmp", "utf8"),
    Buffer.from([0x00, 0x00, 0x00, 0x00, 0x00]),
    Buffer.from(text, "utf8"),
  ]);
  const typeAndData = Buffer.concat([Buffer.from("iTXt", "ascii"), data]);
  const chunk = Buffer.allocUnsafe(12 + data.length);
  chunk.writeUInt32BE(data.length, 0);
  chunk.write("iTXt", 4, "ascii");
  data.copy(chunk, 8);
  chunk.writeUInt32BE(crc32(typeAndData), 8 + data.length);
  return chunk;
}

export function injectXmpIntoBuffer(pngBuffer: Buffer, params: PngXmpParams): Buffer {
  const dimensions = getPngDimensions(pngBuffer);
  if (!dimensions) return pngBuffer;
  const populatedParams: PngXmpParams = {
    ...params,
    width: params.width ?? dimensions.width,
    height: params.height ?? dimensions.height,
  };
  const chunk = buildITxtChunk(buildXmpString(populatedParams));
  const insertAt = 8 + 25;
  return Buffer.concat([pngBuffer.subarray(0, insertAt), chunk, pngBuffer.subarray(insertAt)]);
}

// PNG iTXt chunk injection for Draw Things-compatible XMP metadata.
import os from "node:os";

export type PngAnalysisMetadata = {
  schema: "ceveyne.image-analysis/v1";
  tool: "detect_object" | "annotate_image";
  sourceNotation: string;
  inference?: {
    model: string;
    query?: string;
    detectorPromptSha256?: string;
    maxTokens?: number;
    temperature?: number;
    reused?: boolean;
  };
  render?: {
    palette?: boolean;
    color?: string;
    lineWeight?: number;
    frameAdjust?: number | string;
  };
  detections: Array<{
    label: string;
    bbox: { x1: number; y1: number; x2: number; y2: number };
  }>;
};

export type PngXmpParams = {
  prompt?: string;
  model?: string;
  width?: number;
  height?: number;
  sources?: string[];
  mode?: string;
  generatedBy?: string;
  creatorTool?: string;
  createdAt?: string;
  analysis?: PngAnalysisMetadata;
};

const CRC_TABLE = (() => {
  const table = new Int32Array(256);
  for (let n = 0; n < 256; n++) {
    let crc = n;
    for (let bit = 0; bit < 8; bit++) {
      crc = crc & 1 ? (0xedb88320 ^ (crc >>> 1)) : (crc >>> 1);
    }
    table[n] = crc;
  }
  return table;
})();

function crc32(data: Buffer): number {
  let crc = 0xffffffff;
  for (let index = 0; index < data.length; index++) {
    crc = CRC_TABLE[(crc ^ data[index]!) & 0xff]! ^ (crc >>> 8);
  }
  return (crc ^ 0xffffffff) >>> 0;
}

function normPath(absolutePath: string): string {
  const home = os.homedir();
  return absolutePath.startsWith(home) ? `~${absolutePath.slice(home.length)}` : absolutePath;
}

function escapeXml(value: string): string {
  return value
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;");
}

function getPngDimensions(pngBuffer: Buffer): { width: number; height: number } | undefined {
  if (
    pngBuffer.length < 24 ||
    !pngBuffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a])) ||
    pngBuffer.toString("ascii", 12, 16) !== "IHDR"
  ) {
    return undefined;
  }
  return { width: pngBuffer.readUInt32BE(16), height: pngBuffer.readUInt32BE(20) };
}

function buildXmpString(params: PngXmpParams): string {
  const createdAt = escapeXml(params.createdAt ?? new Date().toISOString());
  const metadata: Record<string, unknown> = {};
  if (params.prompt != null) metadata["c"] = params.prompt;
  if (params.model) metadata["model"] = params.model;
  if (typeof params.width === "number" && typeof params.height === "number") {
    metadata["size"] = `${params.width}x${params.height}`;
  }
  if (params.sources?.length) metadata["sources"] = params.sources.map(normPath);
  if (params.mode) metadata["mode"] = params.mode;
  if (params.analysis) metadata["analysis"] = params.analysis;
  metadata["generated_by"] = params.generatedBy ?? "ceveyne/analyse-image";

  const description: string[] = [];
  if (params.prompt) description.push(params.prompt);
  const details: string[] = [];
  if (typeof params.width === "number" && typeof params.height === "number") details.push(`Size: ${params.width}x${params.height}`);
  if (params.model) details.push(`Model: ${params.model}`);
  if (details.length) description.push(details.join(", "));
  if (params.sources?.length) description.push(`Source: ${params.sources.map(normPath).join(", ")}`);

  return [
    `<x:xmpmeta xmlns:x="adobe:ns:meta/" x:xmptk="XMP Core 6.0.0">`,
    `   <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">`,
    `      <rdf:Description rdf:about=""`,
    `            xmlns:dc="http://purl.org/dc/elements/1.1/"`,
    `            xmlns:xmp="http://ns.adobe.com/xap/1.0/"`,
    `            xmlns:exif="http://ns.adobe.com/exif/1.0/">`,
    `         <dc:description><rdf:Alt><rdf:li xml:lang="x-default">${escapeXml(description.join("\n"))}</rdf:li></rdf:Alt></dc:description>`,
    `         <xmp:CreatorTool>${escapeXml(params.creatorTool ?? params.generatedBy ?? "ceveyne/analyse-image")}</xmp:CreatorTool>`,
    `         <xmp:CreateDate>${createdAt}</xmp:CreateDate>`,
    `         <exif:UserComment><rdf:Alt><rdf:li xml:lang="x-default">${JSON.stringify(metadata)}</rdf:li></rdf:Alt></exif:UserComment>`,
    `      </rdf:Description>`,
    `   </rdf:RDF>`,
    `</x:xmpmeta>`,
  ].join("\n");
}

function buildITxtChunk(text: string): Buffer {
  const data = Buffer.concat([
    Buffer.from("XML:com.adobe.xmp", "utf8"),
    Buffer.from([0x00, 0x00, 0x00, 0x00, 0x00]),
    Buffer.from(text, "utf8"),
  ]);
  const typeAndData = Buffer.concat([Buffer.from("iTXt", "ascii"), data]);
  const chunk = Buffer.allocUnsafe(12 + data.length);
  chunk.writeUInt32BE(data.length, 0);
  chunk.write("iTXt", 4, "ascii");
  data.copy(chunk, 8);
  chunk.writeUInt32BE(crc32(typeAndData), 8 + data.length);
  return chunk;
}

export function injectXmpIntoBuffer(pngBuffer: Buffer, params: PngXmpParams): Buffer {
  const dimensions = getPngDimensions(pngBuffer);
  if (!dimensions) return pngBuffer;
  const populatedParams: PngXmpParams = {
    ...params,
    width: params.width ?? dimensions.width,
    height: params.height ?? dimensions.height,
  };
  const chunk = buildITxtChunk(buildXmpString(populatedParams));
  const insertAt = 8 + 25;
  return Buffer.concat([pngBuffer.subarray(0, insertAt), chunk, pngBuffer.subarray(insertAt)]);
}