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

/** Decoded RGBA image, row-major, 4 bytes per pixel. */
export interface Raster {
  width: number;
  height: number;
  data: Uint8Array;
}

export interface Rgb {
  r: number;
  g: number;
  b: number;
}

export async function loadRaster(path: string): Promise<Raster> {
  const bytes = await readFile(path);
  const extension = extname(path).toLowerCase();

  if (extension === ".png") {
    const png = PNG.sync.read(bytes);
    return { width: png.width, height: png.height, data: new Uint8Array(png.data) };
  }

  if (extension === ".jpg" || extension === ".jpeg") {
    // `useTArray` keeps the result a Uint8Array rather than a Node Buffer.
    const decoded = jpeg.decode(bytes, { useTArray: true });
    return { width: decoded.width, height: decoded.height, data: new Uint8Array(decoded.data) };
  }

  throw new Error(`Unsupported image format: ${extension || path}. Use PNG or JPEG.`);
}

export const pixelIndex = (raster: Raster, x: number, y: number): number =>
  (y * raster.width + x) * 4;

export function pixelAt(raster: Raster, x: number, y: number): Rgb {
  const i = pixelIndex(raster, x, y);
  return { r: raster.data[i]!, g: raster.data[i + 1]!, b: raster.data[i + 2]! };
}

/** Perceptual-ish distance. Cheap, and good enough to separate a flat backdrop. */
export function colorDistance(a: Rgb, b: Rgb): number {
  const dr = a.r - b.r;
  const dg = a.g - b.g;
  const db = a.b - b.b;
  // Green weighted highest, matching luminance sensitivity.
  return Math.sqrt(dr * dr * 0.3 + dg * dg * 0.59 + db * db * 0.11);
}

export const luminance = (color: Rgb): number =>
  (color.r * 0.2126 + color.g * 0.7152 + color.b * 0.0722) / 255;

/** Debug output: write a mask or raster to PNG so a human can eyeball it. */
export function rasterToPng(raster: Raster): Buffer {
  const png = new PNG({ width: raster.width, height: raster.height });
  png.data = Buffer.from(raster.data);
  return PNG.sync.write(png);
}
