import { colorDistance, pixelAt, type Raster, type Rgb } from "./Raster.js";

/**
 * Foreground extraction.
 *
 * The reference prompt asks for a plain uniform backdrop, which makes this
 * tractable without a segmentation model: sample the corners for the backdrop
 * colour, then flood-fill inward from the border. Flood-fill rather than a
 * global colour threshold matters — a global test would also erase any part of
 * the *character* that happens to match the backdrop (grey clothing against
 * grey), whereas a fill only removes background actually connected to the edge.
 */

export interface Mask {
  width: number;
  height: number;
  /** 1 = foreground, 0 = background. */
  data: Uint8Array;
  /** Tight bounds of the foreground. */
  bounds: { minX: number; minY: number; maxX: number; maxY: number };
  coverage: number;
}

export interface SegmentOptions {
  /** Colour distance under which a pixel counts as backdrop. */
  tolerance?: number;
  /** Fraction of the shorter side sampled at each corner. */
  cornerSampleRatio?: number;
}

const DEFAULT_TOLERANCE = 42;

/** Median of the four corner patches, so one odd corner cannot skew it. */
export function estimateBackdrop(raster: Raster, sampleRatio = 0.06): Rgb {
  const size = Math.max(2, Math.floor(Math.min(raster.width, raster.height) * sampleRatio));
  const reds: number[] = [];
  const greens: number[] = [];
  const blues: number[] = [];

  const corners: Array<[number, number]> = [
    [0, 0],
    [raster.width - size, 0],
    [0, raster.height - size],
    [raster.width - size, raster.height - size],
  ];

  for (const [originX, originY] of corners) {
    for (let y = originY; y < originY + size; y++) {
      for (let x = originX; x < originX + size; x++) {
        const pixel = pixelAt(raster, x, y);
        reds.push(pixel.r);
        greens.push(pixel.g);
        blues.push(pixel.b);
      }
    }
  }

  return { r: median(reds), g: median(greens), b: median(blues) };
}

export function segmentForeground(raster: Raster, options: SegmentOptions = {}): Mask {
  const tolerance = options.tolerance ?? DEFAULT_TOLERANCE;
  const backdrop = estimateBackdrop(raster, options.cornerSampleRatio);

  const { width, height } = raster;
  const total = width * height;

  // Start as all-foreground; the fill carves the background away.
  const data = new Uint8Array(total).fill(1);

  // Iterative stack rather than recursion — a 1024x1365 image overflows the
  // call stack immediately with a recursive fill.
  const stack: number[] = [];
  const pushIfBackdrop = (x: number, y: number) => {
    if (x < 0 || y < 0 || x >= width || y >= height) return;
    const index = y * width + x;
    if (data[index] === 0) return;
    if (colorDistance(pixelAt(raster, x, y), backdrop) > tolerance) return;
    data[index] = 0;
    stack.push(index);
  };

  for (let x = 0; x < width; x++) {
    pushIfBackdrop(x, 0);
    pushIfBackdrop(x, height - 1);
  }
  for (let y = 0; y < height; y++) {
    pushIfBackdrop(0, y);
    pushIfBackdrop(width - 1, y);
  }

  while (stack.length > 0) {
    const index = stack.pop()!;
    const x = index % width;
    const y = (index - x) / width;
    pushIfBackdrop(x + 1, y);
    pushIfBackdrop(x - 1, y);
    pushIfBackdrop(x, y + 1);
    pushIfBackdrop(x, y - 1);
  }

  const mask: Mask = {
    width,
    height,
    data,
    bounds: { minX: width, minY: height, maxX: -1, maxY: -1 },
    coverage: 0,
  };

  let filled = 0;
  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      if (data[y * width + x] !== 1) continue;
      filled++;
      if (x < mask.bounds.minX) mask.bounds.minX = x;
      if (x > mask.bounds.maxX) mask.bounds.maxX = x;
      if (y < mask.bounds.minY) mask.bounds.minY = y;
      if (y > mask.bounds.maxY) mask.bounds.maxY = y;
    }
  }

  mask.coverage = filled / total;
  if (mask.bounds.maxX < 0) {
    throw new Error("Segmentation found no foreground — is the background plain and uniform?");
  }
  return mask;
}

/**
 * Removes salt-and-pepper speckle left by JPEG ringing near the silhouette
 * edge, which otherwise adds spurious width to the row profile.
 */
export function denoise(mask: Mask, minNeighbours = 5): Mask {
  const { width, height, data } = mask;
  const out = new Uint8Array(data);

  for (let y = 1; y < height - 1; y++) {
    for (let x = 1; x < width - 1; x++) {
      let neighbours = 0;
      for (let dy = -1; dy <= 1; dy++) {
        for (let dx = -1; dx <= 1; dx++) {
          if (dx === 0 && dy === 0) continue;
          neighbours += data[(y + dy) * width + (x + dx)]!;
        }
      }
      const index = y * width + x;
      if (data[index] === 1 && neighbours < minNeighbours - 2) out[index] = 0;
      else if (data[index] === 0 && neighbours > minNeighbours + 1) out[index] = 1;
    }
  }

  return { ...mask, data: out };
}

export const isForeground = (mask: Mask, x: number, y: number): boolean =>
  x >= 0 && y >= 0 && x < mask.width && y < mask.height && mask.data[y * mask.width + x] === 1;

/** Foreground pixel count and horizontal extent for one scanline. */
export function rowProfile(mask: Mask, y: number): { count: number; minX: number; maxX: number } {
  let count = 0;
  let minX = mask.width;
  let maxX = -1;

  for (let x = 0; x < mask.width; x++) {
    if (mask.data[y * mask.width + x] !== 1) continue;
    count++;
    if (x < minX) minX = x;
    if (x > maxX) maxX = x;
  }
  return { count, minX, maxX };
}

function median(values: number[]): number {
  const sorted = [...values].sort((a, b) => a - b);
  const mid = sorted.length >> 1;
  return sorted.length % 2 === 0
    ? Math.round((sorted[mid - 1]! + sorted[mid]!) / 2)
    : sorted[mid]!;
}
