import { readFile } from "node:fs/promises";
import { inflateSync } from "node:zlib";
import type { NormalizedBox } from "@rakazo/contracts";
import { runMediaTool } from "./media-folder-source-probe.js";

/*
 * Pixel checks for the privacy gate and mark quality. Video frames are decoded
 * by ffmpeg straight to 8-bit luma; mask PNGs from the segmenter go through the
 * decoder below (8-bit gray, gray+alpha, RGB, RGBA; not interlaced, not palette).
 */

export interface LumaFrame {
  width: number;
  height: number;
  /** One luma byte per pixel, row-major. */
  luma: Uint8Array;
}

const PNG_SIGNATURE = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
const CHANNELS: Record<number, number> = { 0: 1, 2: 3, 4: 2, 6: 4 };

function paeth(a: number, b: number, c: number): number {
  const p = a + b - c;
  const pa = Math.abs(p - a);
  const pb = Math.abs(p - b);
  const pc = Math.abs(p - c);
  if (pa <= pb && pa <= pc) return a;
  return pb <= pc ? b : c;
}

/** Decode a PNG to luma (Rec. 601 weights). Throws on anything it cannot read exactly. */
export function decodePngLuma(png: Buffer): LumaFrame {
  if (png.length < 8 || !png.subarray(0, 8).equals(PNG_SIGNATURE)) throw new Error("not a PNG");
  let offset = 8;
  let width = 0;
  let height = 0;
  let channels = 0;
  const data: Buffer[] = [];
  while (offset + 8 <= png.length) {
    const length = png.readUInt32BE(offset);
    const type = png.toString("ascii", offset + 4, offset + 8);
    const body = png.subarray(offset + 8, offset + 8 + length);
    offset += 12 + length;
    if (type === "IHDR") {
      width = body.readUInt32BE(0);
      height = body.readUInt32BE(4);
      const depth = body[8];
      const colorType = body[9] ?? -1;
      const interlace = body[12];
      channels = CHANNELS[colorType] ?? 0;
      if (depth !== 8 || channels === 0 || interlace !== 0)
        throw new Error("unsupported PNG layout");
    } else if (type === "IDAT") {
      data.push(body);
    } else if (type === "IEND") {
      break;
    }
  }
  if (width === 0 || height === 0 || channels === 0) throw new Error("PNG has no image header");
  const raw = inflateSync(Buffer.concat(data));
  const stride = width * channels;
  if (raw.length < (stride + 1) * height) throw new Error("PNG data is truncated");
  const pixels = Buffer.alloc(stride * height);
  for (let y = 0; y < height; y++) {
    const filter = raw[y * (stride + 1)];
    const line = raw.subarray(y * (stride + 1) + 1, (y + 1) * (stride + 1));
    const row = y * stride;
    for (let x = 0; x < stride; x++) {
      const left = x >= channels ? (pixels[row + x - channels] ?? 0) : 0;
      const up = y > 0 ? (pixels[row - stride + x] ?? 0) : 0;
      const upLeft = y > 0 && x >= channels ? (pixels[row - stride + x - channels] ?? 0) : 0;
      const value = line[x] ?? 0;
      let out: number;
      if (filter === 0) out = value;
      else if (filter === 1) out = value + left;
      else if (filter === 2) out = value + up;
      else if (filter === 3) out = value + ((left + up) >> 1);
      else if (filter === 4) out = value + paeth(left, up, upLeft);
      else throw new Error("unknown PNG filter");
      pixels[row + x] = out & 0xff;
    }
  }
  const luma = new Uint8Array(width * height);
  for (let i = 0; i < width * height; i++) {
    const base = i * channels;
    if (channels <= 2) {
      luma[i] = pixels[base] ?? 0;
    } else {
      const r = pixels[base] ?? 0;
      const g = pixels[base + 1] ?? 0;
      const b = pixels[base + 2] ?? 0;
      luma[i] = Math.round(0.299 * r + 0.587 * g + 0.114 * b);
    }
  }
  return { width, height, luma };
}

export async function readPngLuma(file: string): Promise<LumaFrame> {
  return decodePngLuma(await readFile(file));
}

/**
 * Decode the frame shown at `tMs` as luma, scaled to `width` x `height` with
 * nearest-neighbour sampling (a no-op at the display size). Throws unless
 * ffmpeg returns a full frame.
 */
export async function extractLuma(
  file: string,
  input: { tMs: number; width: number; height: number; still?: boolean },
  options: { ffmpegPath?: string; signal?: AbortSignal } = {},
): Promise<LumaFrame> {
  const seek = input.still ? [] : ["-ss", (Math.max(0, input.tMs) / 1000).toFixed(3)];
  const { stdout } = await runMediaTool(
    options.ffmpegPath ?? "ffmpeg",
    [
      "-v",
      "error",
      ...seek,
      "-i",
      file,
      "-frames:v",
      "1",
      "-vf",
      `scale=${input.width}:${input.height}:flags=neighbor`,
      "-f",
      "rawvideo",
      "-pix_fmt",
      "gray",
      "-",
    ],
    { signal: options.signal },
  );
  const size = input.width * input.height;
  if (stdout.length < size) throw new Error(`no frame at ${input.tMs} ms`);
  return {
    width: input.width,
    height: input.height,
    luma: new Uint8Array(stdout.subarray(0, size)),
  };
}

/** Bounding box of the bright (>= 128) pixels of a mask, or null when the mask is empty. */
export function maskBounds(mask: LumaFrame): NormalizedBox | null {
  let minX = mask.width;
  let minY = mask.height;
  let maxX = -1;
  let maxY = -1;
  for (let y = 0; y < mask.height; y++) {
    for (let x = 0; x < mask.width; x++) {
      if ((mask.luma[y * mask.width + x] ?? 0) < 128) continue;
      minX = Math.min(minX, x);
      minY = Math.min(minY, y);
      maxX = Math.max(maxX, x);
      maxY = Math.max(maxY, y);
    }
  }
  if (maxX < 0) return null;
  return {
    x: minX / mask.width,
    y: minY / mask.height,
    w: (maxX + 1 - minX) / mask.width,
    h: (maxY + 1 - minY) / mask.height,
  };
}

/** A luma step at least this large marks a detail pixel (an edge someone could recognise). */
export const DETAIL_MIN_STEP = 12;
/** A detail pixel within this many luma levels of the source counts as unchanged. */
export const UNCHANGED_MAX_DIFF = 2;
/** Looser tolerance for the control area outside treated regions (codec noise). */
const CONTROL_MAX_DIFF = 6;
/** Regions with fewer detail pixels show nothing identifying. */
const MIN_DETAIL_PIXELS = 16;

function pixelRange(box: NormalizedBox, width: number, height: number) {
  return {
    x0: Math.max(0, Math.floor(box.x * width)),
    y0: Math.max(0, Math.floor(box.y * height)),
    x1: Math.min(width, Math.ceil((box.x + box.w) * width)),
    y1: Math.min(height, Math.ceil((box.y + box.h) * height)),
  };
}

function detailStep(frame: LumaFrame, x: number, y: number): number {
  const at = y * frame.width + x;
  const value = frame.luma[at] ?? 0;
  const right = x + 1 < frame.width ? (frame.luma[at + 1] ?? value) : value;
  const down = y + 1 < frame.height ? (frame.luma[at + frame.width] ?? value) : value;
  return Math.abs(value - right) + Math.abs(value - down);
}

/**
 * Residual visibility of a region (0..1): the share of the source's detail
 * pixels (edges) that the treated frame left unchanged, value and edge
 * strength alike. Blur, pixelation, fills, and inpainting all rewrite edges,
 * so a treated region scores near 0, an untouched one 1. A flat region
 * without detail scores 0: there is nothing to recognise.
 *
 * `source` and `treated` are pixel-aligned at full resolution; when `treated`
 * is a crop of the source, `offset` is the crop origin in source pixels and
 * detail outside the crop is invisible. Optional `mask` (normalized over the
 * source frame) limits the region to mask pixels.
 */
export function residualVisibility(input: {
  source: LumaFrame;
  treated: LumaFrame;
  box: NormalizedBox;
  offset?: { x: number; y: number };
  mask?: LumaFrame;
}): number {
  const { source, treated, box, mask } = input;
  const offset = input.offset ?? { x: 0, y: 0 };
  const { x0, y0, x1, y1 } = pixelRange(box, source.width, source.height);
  let detail = 0;
  let unchanged = 0;
  for (let y = y0; y < y1; y++) {
    for (let x = x0; x < x1; x++) {
      if (mask) {
        const mx = Math.min(mask.width - 1, Math.floor((x / source.width) * mask.width));
        const my = Math.min(mask.height - 1, Math.floor((y / source.height) * mask.height));
        if ((mask.luma[my * mask.width + mx] ?? 0) < 128) continue;
      }
      const step = detailStep(source, x, y);
      if (step < DETAIL_MIN_STEP) continue;
      const tx = x - offset.x;
      const ty = y - offset.y;
      if (tx < 0 || ty < 0 || tx >= treated.width || ty >= treated.height) continue;
      detail += 1;
      const value = source.luma[y * source.width + x] ?? 0;
      // Visible means the pixel and its edge survived, not a value that happens to match.
      if (
        Math.abs(value - (treated.luma[ty * treated.width + tx] ?? 0)) <= UNCHANGED_MAX_DIFF &&
        Math.abs(step - detailStep(treated, tx, ty)) <= 2 * UNCHANGED_MAX_DIFF
      )
        unchanged += 1;
    }
  }
  return detail < MIN_DETAIL_PIXELS ? 0 : unchanged / detail;
}

/**
 * Share of detail pixels outside every treated box that the treated frame kept
 * (within codec tolerance). A low value means the two frames are not the same
 * moment or framing, so a residual measured between them proves nothing.
 * Returns 1 when the control area has too little detail to judge.
 */
export function controlAgreement(input: {
  source: LumaFrame;
  treated: LumaFrame;
  boxes: NormalizedBox[];
  offset?: { x: number; y: number };
}): number {
  const { source, treated, boxes } = input;
  const offset = input.offset ?? { x: 0, y: 0 };
  const ranges = boxes.map((box) => pixelRange(box, source.width, source.height));
  let detail = 0;
  let kept = 0;
  for (let y = 0; y < source.height; y++) {
    for (let x = 0; x < source.width; x++) {
      const tx = x - offset.x;
      const ty = y - offset.y;
      if (tx < 0 || ty < 0 || tx >= treated.width || ty >= treated.height) continue;
      if (ranges.some((r) => x >= r.x0 - 2 && x < r.x1 + 2 && y >= r.y0 - 2 && y < r.y1 + 2))
        continue;
      if (detailStep(source, x, y) < DETAIL_MIN_STEP) continue;
      detail += 1;
      const value = source.luma[y * source.width + x] ?? 0;
      if (Math.abs(value - (treated.luma[ty * treated.width + tx] ?? 0)) <= CONTROL_MAX_DIFF)
        kept += 1;
    }
  }
  return detail < MIN_DETAIL_PIXELS ? 1 : kept / detail;
}

/**
 * Temporal stability of a treated region over consecutive frames (0..1):
 * 1 when the region changes no more between neighbours than the rest of the
 * frame, falling to 0 as it flickers 24+ luma levels more on average.
 */
export function temporalStability(frames: LumaFrame[], box: NormalizedBox): number {
  if (frames.length < 2) return 1;
  let excess = 0;
  for (let i = 1; i < frames.length; i++) {
    const a = frames[i - 1]!;
    const b = frames[i]!;
    if (a.width !== b.width || a.height !== b.height) throw new Error("frame sizes differ");
    const r = pixelRange(box, a.width, a.height);
    let inside = 0;
    let insideCount = 0;
    let outside = 0;
    let outsideCount = 0;
    for (let y = 0; y < a.height; y++) {
      for (let x = 0; x < a.width; x++) {
        const at = y * a.width + x;
        const diff = Math.abs((a.luma[at] ?? 0) - (b.luma[at] ?? 0));
        if (x >= r.x0 && x < r.x1 && y >= r.y0 && y < r.y1) {
          inside += diff;
          insideCount += 1;
        } else {
          outside += diff;
          outsideCount += 1;
        }
      }
    }
    const meanInside = insideCount ? inside / insideCount : 0;
    const meanOutside = outsideCount ? outside / outsideCount : 0;
    excess += Math.max(0, meanInside - meanOutside);
  }
  return Math.max(0, Math.min(1, 1 - excess / (frames.length - 1) / 24));
}