import sharp from 'sharp'

function dataUrlToBuffer(dataUrl: string): Buffer {
  const match = dataUrl.match(/^data:image\/png;base64,(.+)$/);
  if (!match) {
    throw new Error('Expected PNG data URL');
  }
  return Buffer.from(match[1], 'base64');
}

export interface FuseWeights {
  u2: number;
  heur: number;
}

export async function fuseMasksPNG(
  u2Mask: string,
  heurMask: string,
  weights: FuseWeights = { u2: 1, heur: 0.8 },
  finalBlurPx = 1,
): Promise<string> {
  const u2Buffer = dataUrlToBuffer(u2Mask);
  const heurBuffer = dataUrlToBuffer(heurMask);

  const u2Sharp = sharp(u2Buffer);
  const heurSharp = sharp(heurBuffer);

  const u2Meta = await u2Sharp.metadata();
  if (!u2Meta.width || !u2Meta.height) {
    throw new Error('u2 mask missing dimensions');
  }

  const width = u2Meta.width;
  const height = u2Meta.height;

  const u2Raw = await u2Sharp.ensureAlpha().raw().toBuffer();
  const heurRaw = await heurSharp.ensureAlpha().resize(width, height, { fit: 'fill' }).raw().toBuffer();

  const output = Buffer.alloc(width * height * 4);
  const weightU2 = Math.max(0, weights.u2 ?? 0);
  const weightHeur = Math.max(0, weights.heur ?? 0);

  for (let i = 0; i < width * height; i += 1) {
    const u2Alpha = u2Raw[i * 4 + 3] / 255;
    const heurAlpha = heurRaw[i * 4 + 3] / 255;
    const mask = 1 - Math.pow(1 - u2Alpha, weightU2) * Math.pow(1 - heurAlpha, weightHeur);
    const alpha = Math.round(Math.max(0, Math.min(1, mask)) * 255);

    output[i * 4 + 0] = 128;
    output[i * 4 + 1] = 128;
    output[i * 4 + 2] = 128;
    output[i * 4 + 3] = alpha;
  }

  let pipeline = sharp(output, { raw: { width, height, channels: 4 } });
  if (finalBlurPx > 0) {
    pipeline = pipeline.blur(finalBlurPx);
  }
  const png = await pipeline.png().toBuffer();
  return `data:image/png;base64,${png.toString('base64')}`;
}

type GrayBytes = { data: Uint8Array; width: number; height: number };

async function readMaskToGrayBytes(mask: string, width?: number, height?: number): Promise<GrayBytes> {
  const input = dataUrlToBuffer(mask);
  let pipeline = sharp(input).ensureAlpha();
  if (width && height) {
    pipeline = pipeline.resize(width, height, { fit: 'fill' });
  }
  const { data, info } = await pipeline.raw().toBuffer({ resolveWithObject: true });
  const w = info.width ?? width ?? 0;
  const h = info.height ?? height ?? 0;
  if (!(w > 0 && h > 0)) throw new Error('mask missing dimensions');

  const rgba = new Uint8Array(data);
  let alphaMin = 255, alphaMax = 0, redMin = 255, redMax = 0;
  for (let i = 0; i < rgba.length; i += 4) {
    const a = rgba[i + 3];
    const r = rgba[i];
    if (a < alphaMin) alphaMin = a;
    if (a > alphaMax) alphaMax = a;
    if (r < redMin) redMin = r;
    if (r > redMax) redMax = r;
  }
  const alphaVaries = alphaMax > alphaMin;
  const redVaries = redMax > redMin;
  const useAlpha = alphaVaries || !redVaries;

  const out = new Uint8Array(w * h);
  for (let i = 0, j = 0; i < out.length; i++, j += 4) {
    out[i] = useAlpha ? rgba[j + 3] : rgba[j];
  }
  return { data: out, width: w, height: h };
}

async function encodeGrayBytesPng(bytes: Uint8Array, width: number, height: number): Promise<string> {
  const png = await sharp(Buffer.from(bytes), { raw: { width, height, channels: 1 } }).png().toBuffer();
  return `data:image/png;base64,${png.toString('base64')}`;
}

export async function unionMasksPNG(
  maskA: string,
  maskB: string,
  opts?: { threshold?: number },
): Promise<string> {
  const threshold = Math.max(1, Math.min(254, Math.round(opts?.threshold ?? 32)));
  const a = await readMaskToGrayBytes(maskA);
  const b = await readMaskToGrayBytes(maskB, a.width, a.height);
  const out = new Uint8Array(a.data.length);
  for (let i = 0; i < out.length; i++) {
    const va = a.data[i] >= threshold ? 255 : 0;
    const vb = b.data[i] >= threshold ? 255 : 0;
    out[i] = va || vb ? 255 : 0;
  }
  return encodeGrayBytesPng(out, a.width, a.height);
}

export async function dilateMaskPNG(
  mask: string,
  radius = 4,
  opts?: { threshold?: number },
): Promise<string> {
  const r = Math.max(0, Math.round(radius));
  if (!r) return mask;
  const threshold = Math.max(1, Math.min(254, Math.round(opts?.threshold ?? 8)));
  const { data, width, height } = await readMaskToGrayBytes(mask);
  const blurred = await sharp(Buffer.from(data), { raw: { width, height, channels: 1 } })
    .blur(r)
    .threshold(threshold)
    .png()
    .toBuffer();
  return `data:image/png;base64,${blurred.toString('base64')}`;
}