import type { BoundingBox, Point, RegionMask, RGBColor } from '@/types/mosaic'
import { colorDistance, relativeLuminance } from './colors'
import { createRegionMask, summarizeRegionMask } from './regionMask'

export interface FloodFillOptions {
  tolerance: number
  darkBoundaryThreshold: number
  connectivity: 4 | 8
  occupancy?: Uint8Array
}

export interface FloodFillResult {
  mask: RegionMask
  pixelCount: number
  boundingBox: BoundingBox
  labelPosition: Point
  color: RGBColor
}

const pixelColor = (data: Uint8ClampedArray, index: number): RGBColor => ({
  r: data[index * 4] ?? 0,
  g: data[index * 4 + 1] ?? 0,
  b: data[index * 4 + 2] ?? 0,
})

const median = (values: number[]): number => {
  values.sort((first, second) => first - second)
  return values[Math.floor(values.length / 2)] ?? 0
}

export function getFloodFillReferenceColor(
  image: ImageData,
  startX: number,
  startY: number,
  tolerance: number,
): RGBColor {
  const { width, height, data } = image
  const startIndex = startY * width + startX
  if (tolerance !== 1) return pixelColor(data, startIndex)

  const reds: number[] = []
  const greens: number[] = []
  const blues: number[] = []
  for (let y = Math.max(0, startY - 2); y <= Math.min(height - 1, startY + 2); y += 1) {
    for (let x = Math.max(0, startX - 2); x <= Math.min(width - 1, startX + 2); x += 1) {
      const color = pixelColor(data, y * width + x)
      reds.push(color.r)
      greens.push(color.g)
      blues.push(color.b)
    }
  }
  return { r: median(reds), g: median(greens), b: median(blues) }
}

export function floodFillRegion(
  image: ImageData,
  startX: number,
  startY: number,
  options: FloodFillOptions,
): FloodFillResult | null {
  const { width, height, data } = image
  if (startX < 0 || startY < 0 || startX >= width || startY >= height) return null
  const startIndex = startY * width + startX
  if (options.occupancy?.[startIndex]) return null

  const referenceColor = getFloodFillReferenceColor(image, startX, startY, options.tolerance)
  const referenceIsDark = relativeLuminance(referenceColor) <= options.darkBoundaryThreshold
  const visited = new Uint8Array(width * height)
  const stack: number[] = [startIndex]
  const runs: number[] = []
  const reds: number[] = []
  const greens: number[] = []
  const blues: number[] = []
  let pixelCount = 0

  const isAccepted = (index: number) => {
    if (visited[index] || options.occupancy?.[index]) return false
    if (index === startIndex) return true
    const color = pixelColor(data, index)
    if (!referenceIsDark && relativeLuminance(color) <= options.darkBoundaryThreshold) return false
    return colorDistance(color, referenceColor) <= options.tolerance
  }

  const sample = (index: number) => {
    const color = pixelColor(data, index)
    if (reds.length < 8000) {
      reds.push(color.r)
      greens.push(color.g)
      blues.push(color.b)
      return
    }
    const replacement = Math.floor(Math.random() * pixelCount)
    if (replacement < 8000) {
      reds[replacement] = color.r
      greens[replacement] = color.g
      blues[replacement] = color.b
    }
  }

  while (stack.length) {
    const index = stack.pop()
    if (index === undefined || !isAccepted(index)) continue
    const y = Math.floor(index / width)
    let left = index % width
    let right = left
    while (left > 0 && isAccepted(y * width + left - 1)) left -= 1
    while (right + 1 < width && isAccepted(y * width + right + 1)) right += 1

    for (let x = left; x <= right; x += 1) {
      const pixelIndex = y * width + x
      visited[pixelIndex] = 1
      pixelCount += 1
      // Avoid both ends of the scanline, where outlines and JPEG artifacts are most likely.
      if (x >= left + 2 && x <= right - 2) sample(pixelIndex)
    }
    runs.push(y, left, right)

    for (const neighborY of [y - 1, y + 1]) {
      if (neighborY < 0 || neighborY >= height) continue
      const extra = options.connectivity === 8 ? 1 : 0
      const scanStart = Math.max(0, left - extra)
      const scanEnd = Math.min(width - 1, right + extra)
      let insideCandidate = false
      for (let x = scanStart; x <= scanEnd; x += 1) {
        const neighborIndex = neighborY * width + x
        const candidate = isAccepted(neighborIndex)
        if (candidate && !insideCandidate) stack.push(neighborIndex)
        insideCandidate = candidate
      }
    }
  }
  if (!pixelCount) return null
  if (!reds.length) {
    reds.push(referenceColor.r)
    greens.push(referenceColor.g)
    blues.push(referenceColor.b)
  }

  const orderedRuns: Array<[number, number, number]> = []
  for (let index = 0; index < runs.length; index += 3) {
    orderedRuns.push([runs[index] ?? 0, runs[index + 1] ?? 0, runs[index + 2] ?? 0])
  }
  orderedRuns.sort((first, second) => first[0] - second[0] || first[1] - second[1])
  const mask = createRegionMask(width, height, orderedRuns.flatMap((run) => run))
  const summary = summarizeRegionMask(mask)
  if (!summary) return null

  return {
    mask,
    pixelCount,
    boundingBox: summary.boundingBox,
    labelPosition: summary.labelPosition,
    color: { r: median(reds), g: median(greens), b: median(blues) },
  }
}
