import fs from "node:fs/promises";
import path from "node:path";
import jpeg from "jpeg-js";
import pixelmatch from "pixelmatch";
import { PNG } from "pngjs";

const root = path.resolve(process.argv[2] ?? "docs/visual-parity");
const maxDiffRatio = Number(process.env.VISUAL_MAX_DIFF_RATIO ?? "0.003");
const threshold = Number(process.env.VISUAL_PIXEL_THRESHOLD ?? "0.1");
const writeDiffs = process.env.VISUAL_WRITE_DIFFS === "1";

async function walk(directory) {
  const entries = await fs.readdir(directory, { withFileTypes: true });
  const files = await Promise.all(
    entries.map((entry) => {
      const target = path.join(directory, entry.name);
      return entry.isDirectory() ? walk(target) : [target];
    }),
  );
  return files.flat();
}

const files = await walk(root);
const sources = files.filter(
  (file) => path.basename(file).startsWith("source-") && /\.(?:jpe?g|png)$/i.test(file),
);

if (sources.length === 0) {
  throw new Error(`No source baseline images found under ${root}`);
}

let failed = false;
const results = [];

for (const sourcePath of sources.sort()) {
  const portPath = path.join(path.dirname(sourcePath), path.basename(sourcePath).replace(/^source-/, "port-"));
  const [sourceBytes, portBytes] = await Promise.all([fs.readFile(sourcePath), fs.readFile(portPath)]);
  const decode = (bytes) =>
    bytes[0] === 0xff && bytes[1] === 0xd8
      ? jpeg.decode(bytes, { useTArray: true })
      : PNG.sync.read(bytes);
  const source = decode(sourceBytes);
  const port = decode(portBytes);

  if (source.width !== port.width || source.height !== port.height) {
    failed = true;
    results.push({
      pair: path.relative(root, sourcePath).replace(/^source-/, ""),
      source: `${source.width}x${source.height}`,
      port: `${port.width}x${port.height}`,
      error: "dimension mismatch",
    });
    continue;
  }

  const pixelCount = source.width * source.height;
  const diff = writeDiffs ? new PNG({ width: source.width, height: source.height }) : null;
  const differentPixels = pixelmatch(source.data, port.data, diff?.data ?? null, source.width, source.height, {
    threshold,
    includeAA: false,
  });
  const diffRatio = differentPixels / pixelCount;
  if (writeDiffs && differentPixels > 0 && diff) {
    const diffPath = path.join(
      path.dirname(sourcePath),
      path.basename(sourcePath).replace(/^source-/, "diff-").replace(/\.jpe?g$/i, ".png"),
    );
    await fs.writeFile(diffPath, PNG.sync.write(diff));
  }
  if (diffRatio > maxDiffRatio) failed = true;
  results.push({
    pair: path.relative(root, sourcePath).replace(/^source-/, ""),
    differentPixels,
    pixelCount,
    diffRatio: Number(diffRatio.toFixed(8)),
    pass: diffRatio <= maxDiffRatio,
  });
}

console.table(results);
console.log(`Compared ${results.length} source/port screenshot pairs (max diff ratio ${maxDiffRatio}).`);

if (failed) process.exitCode = 1;