import { inflateSync } from "node:zlib";
import type { Bitmap } from "../../src/lib/ocrPreprocess.js";

interface PngChunk {
  type: string;
  data: Buffer;
}

function readChunks(buffer: Buffer): PngChunk[] {
  const signature = buffer.subarray(0, 8);
  if (!signature.equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) {
    throw new Error("Invalid PNG signature.");
  }
  const chunks: PngChunk[] = [];
  let offset = 8;
  while (offset + 12 <= buffer.length) {
    const length = buffer.readUInt32BE(offset);
    const type = buffer.toString("ascii", offset + 4, offset + 8);
    const dataStart = offset + 8;
    const dataEnd = dataStart + length;
    if (dataEnd + 4 > buffer.length) throw new Error("Invalid PNG chunk length.");
    chunks.push({ type, data: buffer.subarray(dataStart, dataEnd) });
    offset = dataEnd + 4;
    if (type === "IEND") break;
  }
  return chunks;
}

function paethPredictor(left: number, up: number, upperLeft: number) {
  const estimate = left + up - upperLeft;
  const leftDistance = Math.abs(estimate - left);
  const upDistance = Math.abs(estimate - up);
  const upperLeftDistance = Math.abs(estimate - upperLeft);
  if (leftDistance <= upDistance && leftDistance <= upperLeftDistance) return left;
  if (upDistance <= upperLeftDistance) return up;
  return upperLeft;
}

function unfilterScanlines(raw: Buffer, width: number, height: number, bytesPerPixel: number) {
  const stride = width * bytesPerPixel;
  const output = Buffer.alloc(stride * height);
  let rawOffset = 0;
  for (let row = 0; row < height; row++) {
    const filter = raw[rawOffset++];
    const rowOffset = row * stride;
    const previousRowOffset = rowOffset - stride;
    for (let col = 0; col < stride; col++) {
      const value = raw[rawOffset++];
      const left = col >= bytesPerPixel ? output[rowOffset + col - bytesPerPixel] : 0;
      const up = row > 0 ? output[previousRowOffset + col] : 0;
      const upperLeft = row > 0 && col >= bytesPerPixel ? output[previousRowOffset + col - bytesPerPixel] : 0;
      let restored = value;
      if (filter === 1) restored = value + left;
      else if (filter === 2) restored = value + up;
      else if (filter === 3) restored = value + Math.floor((left + up) / 2);
      else if (filter === 4) restored = value + paethPredictor(left, up, upperLeft);
      else if (filter !== 0) throw new Error(`Unsupported PNG filter: ${filter}`);
      output[rowOffset + col] = restored & 0xff;
    }
  }
  return output;
}

export function pngBufferToBitmap(buffer: Buffer): Bitmap {
  const chunks = readChunks(buffer);
  const ihdr = chunks.find((chunk) => chunk.type === "IHDR")?.data;
  if (!ihdr) throw new Error("PNG missing IHDR.");
  const width = ihdr.readUInt32BE(0);
  const height = ihdr.readUInt32BE(4);
  const bitDepth = ihdr[8];
  const colorType = ihdr[9];
  const compression = ihdr[10];
  const filter = ihdr[11];
  const interlace = ihdr[12];
  if (bitDepth !== 8 || compression !== 0 || filter !== 0 || interlace !== 0) {
    throw new Error("Unsupported PNG format.");
  }
  const sourceBytesPerPixel = colorType === 6 ? 4 : colorType === 2 ? 3 : 0;
  if (!sourceBytesPerPixel) throw new Error(`Unsupported PNG color type: ${colorType}`);
  const idat = Buffer.concat(chunks.filter((chunk) => chunk.type === "IDAT").map((chunk) => chunk.data));
  const unfiltered = unfilterScanlines(inflateSync(idat), width, height, sourceBytesPerPixel);
  if (colorType === 6) return { data: unfiltered, width, height };

  const rgba = Buffer.alloc(width * height * 4);
  for (let pixel = 0; pixel < width * height; pixel++) {
    rgba[pixel * 4] = unfiltered[pixel * 3];
    rgba[pixel * 4 + 1] = unfiltered[pixel * 3 + 1];
    rgba[pixel * 4 + 2] = unfiltered[pixel * 3 + 2];
    rgba[pixel * 4 + 3] = 255;
  }
  return { data: rgba, width, height };
}