import { blurMaskValues, despeckleMaskValues, dilateMaskValues, erodeMaskValues } from "./maskRaster";

type ChromaKeySettings = { color: string; tolerance: number; softness: number; feather: number; choke: number; despeckle: number; spill: number };

export async function createChromaKeyPreview(source: string, width: number, height: number, settings: ChromaKeySettings) {
  return render(source, width, height, settings, false);
}

export async function createChromaKeyMask(source: string, width: number, height: number, settings: ChromaKeySettings) {
  return render(source, width, height, settings, true);
}

async function render(source: string, width: number, height: number, settings: ChromaKeySettings, maskOnly: boolean) {
  const canvas = document.createElement("canvas");
  canvas.width = Math.max(1, Math.round(width));
  canvas.height = Math.max(1, Math.round(height));
  const context = canvas.getContext("2d");
  if (!context) return source;
  context.drawImage(await loadImage(source), 0, 0, canvas.width, canvas.height);
  const data = context.getImageData(0, 0, canvas.width, canvas.height);
  const alpha = chromaKeyAlpha(data, canvas.width, canvas.height, settings);
  for (let pixel = 0; pixel < alpha.length; pixel++) {
    const index = pixel * 4;
    if (maskOnly) data.data[index] = data.data[index + 1] = data.data[index + 2] = 255;
    data.data[index + 3] = alpha[pixel] ?? 255;
  }
  context.putImageData(data, 0, 0);
  return canvas.toDataURL("image/png");
}

function chromaKeyAlpha(data: ImageData, width: number, height: number, settings: ChromaKeySettings) {
  const hex = settings.color.replace("#", "");
  const key = { r: Number.parseInt(hex.slice(0, 2), 16), g: Number.parseInt(hex.slice(2, 4), 16), b: Number.parseInt(hex.slice(4, 6), 16) };
  const alpha = new Uint8ClampedArray(width * height);
  for (let pixel = 0; pixel < alpha.length; pixel++) {
    const index = pixel * 4;
    const red = data.data[index] ?? 0;
    const green = data.data[index + 1] ?? 0;
    const blue = data.data[index + 2] ?? 0;
    const distance = Math.hypot(red - key.r, green - key.g, blue - key.b);
    const tolerance = Math.max(0, Math.min(255, settings.tolerance));
    const softness = Math.max(0, Math.min(255, settings.softness));
    const edgeKeep = distance <= tolerance ? 0 : softness > 0 && distance < tolerance + softness ? (distance - tolerance) / softness : 1;
    const dominant = key.g >= key.r && key.g >= key.b ? green : key.r >= key.b ? red : blue;
    const neutral = key.g >= key.r && key.g >= key.b ? Math.max(red, blue) : key.r >= key.b ? Math.max(green, blue) : Math.max(red, green);
    const spillKeep = 1 - Math.max(0, dominant - neutral) / 255 * Math.max(0, Math.min(100, settings.spill)) / 100;
    alpha[pixel] = Math.round((data.data[index + 3] ?? 255) * Math.max(0, Math.min(edgeKeep, spillKeep)));
  }
  let next: Uint8ClampedArray<ArrayBufferLike> = alpha;
  const despeckle = Math.round(Math.max(0, Math.min(20, settings.despeckle)));
  const choke = Math.round(Math.max(-20, Math.min(20, settings.choke)));
  const feather = Math.round(Math.max(0, Math.min(20, settings.feather)));
  if (despeckle > 0) next = despeckleMaskValues(next, width, height, despeckle);
  if (choke > 0) next = erodeMaskValues(next, width, height, choke);
  if (choke < 0) next = dilateMaskValues(next, width, height, -choke);
  return feather > 0 ? blurMaskValues(next, width, height, feather) : next;
}

function loadImage(source: string) {
  return new Promise<HTMLImageElement>((resolve, reject) => {
    const image = new Image();
    image.onload = () => resolve(image);
    image.onerror = () => reject(new Error("Failed to load image"));
    image.src = source;
  });
}