export type MutableRgbaBuffer = Uint8Array | Uint8ClampedArray;

const SOFT_EDGE_ALPHA_THRESHOLD = 224;
const FOREGROUND_NEIGHBOR_ALPHA_THRESHOLD = 96;

function clamp01(value: number) {
  return Math.max(0, Math.min(1, value));
}

function lerp(from: number, to: number, t: number) {
  return from + (to - from) * clamp01(t);
}

function computeGreenBackgroundScore(
  red: number,
  green: number,
  blue: number,
  alpha: number,
) {
  if (alpha === 0) {
    return 1;
  }

  const greenLead = green - Math.max(red, blue);
  if (green < 52 || greenLead <= 8) {
    return 0;
  }

  const greenRatio = green / Math.max(1, red + blue);
  if (greenRatio <= 0.52) {
    return 0;
  }

  return clamp01(
    ((green - 52) / 168) * 0.22 +
      ((greenLead - 8) / 96) * 0.53 +
      ((greenRatio - 0.52) / 0.82) * 0.25,
  );
}

function computeWhiteBackgroundScore(
  red: number,
  green: number,
  blue: number,
  alpha: number,
) {
  if (alpha === 0) {
    return 1;
  }

  const maxChannel = Math.max(red, green, blue);
  const minChannel = Math.min(red, green, blue);
  const average = (red + green + blue) / 3;
  if (average < 188 || minChannel < 168) {
    return 0;
  }

  const spread = maxChannel - minChannel;
  const neutrality = 1 - clamp01((spread - 6) / 34);
  const brightness = clamp01((average - 188) / 55);
  const floor = clamp01((minChannel - 168) / 60);
  return clamp01(neutrality * (brightness * 0.85 + floor * 0.15));
}

function collectForegroundNeighborColor(
  pixels: MutableRgbaBuffer,
  width: number,
  height: number,
  x: number,
  y: number,
  backgroundMask: Uint8Array,
  backgroundHints: Float32Array,
) {
  let totalWeight = 0;
  let totalRed = 0;
  let totalGreen = 0;
  let totalBlue = 0;

  for (let offsetY = -2; offsetY <= 2; offsetY += 1) {
    for (let offsetX = -2; offsetX <= 2; offsetX += 1) {
      if (offsetX === 0 && offsetY === 0) {
        continue;
      }

      const nextX = x + offsetX;
      const nextY = y + offsetY;
      if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
        continue;
      }

      const nextPixelIndex = nextY * width + nextX;
      if (backgroundMask[nextPixelIndex]) {
        continue;
      }

      if ((backgroundHints[nextPixelIndex] ?? 0) >= 0.18) {
        continue;
      }

      const nextOffset = nextPixelIndex * 4;
      const nextAlpha = pixels[nextOffset + 3] ?? 0;
      if (nextAlpha < FOREGROUND_NEIGHBOR_ALPHA_THRESHOLD) {
        continue;
      }

      const distance = Math.abs(offsetX) + Math.abs(offsetY);
      const weight =
        (nextAlpha / 255) *
        (distance <= 1 ? 1.8 : distance === 2 ? 1.2 : 0.7);

      totalWeight += weight;
      totalRed += (pixels[nextOffset] ?? 0) * weight;
      totalGreen += (pixels[nextOffset + 1] ?? 0) * weight;
      totalBlue += (pixels[nextOffset + 2] ?? 0) * weight;
    }
  }

  if (totalWeight <= 0) {
    return null;
  }

  return {
    red: Math.round(totalRed / totalWeight),
    green: Math.round(totalGreen / totalWeight),
    blue: Math.round(totalBlue / totalWeight),
  };
}

export function removeBackgroundFromRgba(
  pixels: MutableRgbaBuffer,
  width: number,
  height: number,
) {
  const pixelCount = width * height;
  if (pixelCount <= 0) {
    return false;
  }

  const backgroundMask = new Uint8Array(pixelCount);
  const greenScores = new Float32Array(pixelCount);
  const whiteScores = new Float32Array(pixelCount);
  const backgroundHints = new Float32Array(pixelCount);
  const queue: number[] = [];
  let queueIndex = 0;
  let changed = false;

  for (let pixelIndex = 0; pixelIndex < pixelCount; pixelIndex += 1) {
    const offset = pixelIndex * 4;
    const red = pixels[offset] ?? 0;
    const green = pixels[offset + 1] ?? 0;
    const blue = pixels[offset + 2] ?? 0;
    const alpha = pixels[offset + 3] ?? 0;
    const greenScore = computeGreenBackgroundScore(red, green, blue, alpha);
    const whiteScore = computeWhiteBackgroundScore(red, green, blue, alpha);
    const transparencyHint = clamp01((56 - alpha) / 56) * 0.75;

    greenScores[pixelIndex] = greenScore;
    whiteScores[pixelIndex] = whiteScore;
    backgroundHints[pixelIndex] = Math.max(
      greenScore,
      whiteScore,
      transparencyHint,
    );
  }

  const trySeedBackground = (pixelIndex: number) => {
    if (backgroundMask[pixelIndex]) {
      return;
    }

    const offset = pixelIndex * 4;
    const alpha = pixels[offset + 3] ?? 0;
    const strongCandidate =
      alpha < 40 ||
      (greenScores[pixelIndex] ?? 0) > 0.12 ||
      (whiteScores[pixelIndex] ?? 0) > 0.32;

    if (!strongCandidate) {
      return;
    }

    backgroundMask[pixelIndex] = 1;
    queue.push(pixelIndex);
  };

  for (let x = 0; x < width; x += 1) {
    trySeedBackground(x);
    trySeedBackground((height - 1) * width + x);
  }

  for (let y = 1; y < height - 1; y += 1) {
    trySeedBackground(y * width);
    trySeedBackground(y * width + width - 1);
  }

  while (queueIndex < queue.length) {
    const pixelIndex = queue[queueIndex]!;
    queueIndex += 1;

    const x = pixelIndex % width;
    const y = Math.floor(pixelIndex / width);
    const neighborIndexes = [
      x > 0 ? pixelIndex - 1 : -1,
      x + 1 < width ? pixelIndex + 1 : -1,
      y > 0 ? pixelIndex - width : -1,
      y + 1 < height ? pixelIndex + width : -1,
    ];

    for (const nextPixelIndex of neighborIndexes) {
      if (nextPixelIndex < 0 || backgroundMask[nextPixelIndex]) {
        continue;
      }

      const nextOffset = nextPixelIndex * 4;
      const nextAlpha = pixels[nextOffset + 3] ?? 0;
      const nextGreenScore = greenScores[nextPixelIndex] ?? 0;
      const nextWhiteScore = whiteScores[nextPixelIndex] ?? 0;
      const nextHint = backgroundHints[nextPixelIndex] ?? 0;
      const reachableSoftEdge =
        nextHint > 0.08 &&
        nextAlpha < SOFT_EDGE_ALPHA_THRESHOLD &&
        (nextGreenScore > 0.04 || nextWhiteScore > 0.08 || nextAlpha < 180);

      if (
        nextAlpha < 40 ||
        nextGreenScore > 0.12 ||
        nextWhiteScore > 0.32 ||
        reachableSoftEdge
      ) {
        backgroundMask[nextPixelIndex] = 1;
        queue.push(nextPixelIndex);
      }
    }
  }

  for (let iteration = 0; iteration < 2; iteration += 1) {
    const expandedMask = new Uint8Array(backgroundMask);

    for (let y = 0; y < height; y += 1) {
      for (let x = 0; x < width; x += 1) {
        const pixelIndex = y * width + x;
        if (expandedMask[pixelIndex]) {
          continue;
        }

        const alpha = pixels[pixelIndex * 4 + 3] ?? 0;
        const hint = backgroundHints[pixelIndex] ?? 0;
        if (alpha >= SOFT_EDGE_ALPHA_THRESHOLD || hint <= 0.06) {
          continue;
        }

        let adjacentBackgroundCount = 0;
        for (let offsetY = -1; offsetY <= 1; offsetY += 1) {
          for (let offsetX = -1; offsetX <= 1; offsetX += 1) {
            if (offsetX === 0 && offsetY === 0) {
              continue;
            }

            const nextX = x + offsetX;
            const nextY = y + offsetY;
            if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
              continue;
            }

            if (backgroundMask[nextY * width + nextX]) {
              adjacentBackgroundCount += 1;
            }
          }
        }

        if (
          adjacentBackgroundCount >= 2 ||
          (adjacentBackgroundCount >= 1 && hint > 0.18)
        ) {
          expandedMask[pixelIndex] = 1;
        }
      }
    }

    backgroundMask.set(expandedMask);
  }

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const pixelIndex = y * width + x;
      if (!backgroundMask[pixelIndex]) {
        continue;
      }

      const offset = pixelIndex * 4;
      const alpha = pixels[offset + 3] ?? 0;
      if (alpha === 0) {
        continue;
      }

      const matteScore = Math.max(
        backgroundHints[pixelIndex] ?? 0,
        greenScores[pixelIndex] ?? 0,
        whiteScores[pixelIndex] ?? 0,
      );

      let foregroundSupport = 0;
      for (let offsetY = -1; offsetY <= 1; offsetY += 1) {
        for (let offsetX = -1; offsetX <= 1; offsetX += 1) {
          if (offsetX === 0 && offsetY === 0) {
            continue;
          }

          const nextX = x + offsetX;
          const nextY = y + offsetY;
          if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
            continue;
          }

          const nextPixelIndex = nextY * width + nextX;
          if (backgroundMask[nextPixelIndex]) {
            continue;
          }

          const nextAlpha = pixels[nextPixelIndex * 4 + 3] ?? 0;
          if (nextAlpha >= FOREGROUND_NEIGHBOR_ALPHA_THRESHOLD) {
            foregroundSupport += 1;
          }
        }
      }

      let nextAlpha = alpha;
      if (matteScore > 0.9 || foregroundSupport === 0) {
        nextAlpha = 0;
      } else if (matteScore > 0.72 && foregroundSupport <= 1) {
        nextAlpha = Math.min(alpha, Math.round(alpha * 0.08));
      } else {
        nextAlpha = Math.min(
          alpha,
          Math.round(alpha * Math.max(0.08, 1 - matteScore * 0.95)),
        );
      }

      if (foregroundSupport >= 3 && matteScore < 0.55) {
        nextAlpha = Math.max(nextAlpha, Math.round(alpha * 0.22));
      }

      if (nextAlpha < 10) {
        nextAlpha = 0;
      }

      if (nextAlpha !== alpha) {
        pixels[offset + 3] = nextAlpha;
        changed = true;
      }
    }
  }

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const pixelIndex = y * width + x;
      const offset = pixelIndex * 4;
      const alpha = pixels[offset + 3] ?? 0;
      if (alpha === 0) {
        continue;
      }

      let touchesTransparentEdge = false;
      for (let offsetY = -1; offsetY <= 1; offsetY += 1) {
        for (let offsetX = -1; offsetX <= 1; offsetX += 1) {
          if (offsetX === 0 && offsetY === 0) {
            continue;
          }

          const nextX = x + offsetX;
          const nextY = y + offsetY;
          if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
            touchesTransparentEdge = true;
            continue;
          }

          const nextPixelIndex = nextY * width + nextX;
          if (
            backgroundMask[nextPixelIndex] ||
            (pixels[nextPixelIndex * 4 + 3] ?? 0) < 16
          ) {
            touchesTransparentEdge = true;
          }
        }
      }

      if (!touchesTransparentEdge) {
        continue;
      }

      const greenScore = greenScores[pixelIndex] ?? 0;
      const whiteScore = whiteScores[pixelIndex] ?? 0;
      const contamination = Math.max(
        greenScore,
        whiteScore,
        backgroundMask[pixelIndex] ? 0.35 : 0,
        alpha < 220 ? ((220 - alpha) / 220) * 0.25 : 0,
      );

      if (contamination < 0.06) {
        continue;
      }

      let red = pixels[offset] ?? 0;
      let green = pixels[offset + 1] ?? 0;
      let blue = pixels[offset + 2] ?? 0;
      const sample = collectForegroundNeighborColor(
        pixels,
        width,
        height,
        x,
        y,
        backgroundMask,
        backgroundHints,
      );
      const blend = clamp01(
        Math.max(contamination * 0.82, touchesTransparentEdge ? 0.22 : 0),
      );

      if (sample) {
        red = Math.round(lerp(red, sample.red, blend));
        green = Math.round(lerp(green, sample.green, blend));
        blue = Math.round(lerp(blue, sample.blue, blend));

        if (greenScore > 0.04) {
          green = Math.min(green, sample.green + 18);
        }

        if (whiteScore > 0.1) {
          red = Math.min(red, sample.red + 26);
          green = Math.min(green, sample.green + 26);
          blue = Math.min(blue, sample.blue + 26);
        }
      } else {
        if (greenScore > 0.04) {
          green = Math.max(
            Math.max(red, blue),
            Math.round(green - (green - Math.max(red, blue)) * 0.78),
          );
        }

        if (whiteScore > 0.12) {
          const spread = Math.max(red, green, blue) - Math.min(red, green, blue);
          if (spread < 20) {
            const tonedValue = Math.round(((red + green + blue) / 3) * 0.88);
            red = Math.min(red, tonedValue);
            green = Math.min(green, tonedValue);
            blue = Math.min(blue, tonedValue);
          }
        }
      }

      let nextAlpha = alpha;
      const edgeFade = Math.max(greenScore * 0.35, whiteScore * 0.28);
      if (edgeFade > 0.08) {
        nextAlpha = Math.min(alpha, Math.round(alpha * (1 - edgeFade)));
        if (nextAlpha < 10) {
          nextAlpha = 0;
        }
      }

      if (
        red !== (pixels[offset] ?? 0) ||
        green !== (pixels[offset + 1] ?? 0) ||
        blue !== (pixels[offset + 2] ?? 0) ||
        nextAlpha !== alpha
      ) {
        pixels[offset] = red;
        pixels[offset + 1] = green;
        pixels[offset + 2] = blue;
        pixels[offset + 3] = nextAlpha;
        changed = true;
      }
    }
  }

  return changed;
}