import { useCallback, useMemo } from 'react'
import { useCanvas } from '../hooks/useCanvas'
import { AXIS, FONT, FRAME, formatTick, niceTicks } from '../plot'
import type { ImageData2D } from '../types'

type Props = {
  image: ImageData2D
  /** maps a data value to rgb */
  color: (v: number) => [number, number, number]
  label: string
  labelColor: string
  /** time range to outline, e.g. the waveform view */
  highlight?: [number, number]
  className?: string
}

const M = { l: 48, r: 12, t: 8, b: 26 }

/** Draws a time-frequency image (spectrogram, phase field) with axes. */
export function ImagePlot({ image, color, label, labelColor, highlight, className }: Props) {
  // Rasterize once per image, then scale to the canvas on every draw.
  const raster = useMemo(() => {
    const { rows, cols, data } = image
    const off = document.createElement('canvas')
    off.width = cols
    off.height = rows
    const ictx = off.getContext('2d')
    if (!ictx) return null
    const img = ictx.createImageData(cols, rows)
    const d = img.data
    for (let r = 0; r < rows; r++) {
      const yImg = rows - 1 - r // row 0 (0 Hz) at the bottom
      for (let c = 0; c < cols; c++) {
        const [R, G, B] = color(data[r * cols + c])
        const o = 4 * (yImg * cols + c)
        d[o] = R
        d[o + 1] = G
        d[o + 2] = B
        d[o + 3] = 255
      }
    }
    ictx.putImageData(img, 0, 0)
    return off
  }, [image, color])

  const draw = useCallback(
    (ctx: CanvasRenderingContext2D, w: number, h: number) => {
      const pw = w - M.l - M.r
      const ph = h - M.t - M.b
      ctx.font = FONT
      if (raster) {
        ctx.imageSmoothingEnabled = false
        ctx.drawImage(raster, M.l, M.t, pw, ph)
      }
      ctx.strokeStyle = FRAME
      ctx.lineWidth = 1
      ctx.strokeRect(M.l + 0.5, M.t + 0.5, pw, ph)
      if (highlight) {
        const x0 = M.l + (Math.max(0, highlight[0]) / image.durationSec) * pw
        const x1 = M.l + (Math.min(image.durationSec, highlight[1]) / image.durationSec) * pw
        ctx.strokeStyle = 'rgba(255, 255, 255, 0.85)'
        ctx.setLineDash([4, 3])
        ctx.strokeRect(Math.round(x0) + 0.5, M.t + 0.5, Math.max(2, Math.round(x1 - x0)), ph)
        ctx.setLineDash([])
      }

      ctx.fillStyle = AXIS
      ctx.textAlign = 'right'
      ctx.textBaseline = 'middle'
      for (const v of niceTicks(0, image.fmax, 4)) {
        if (v > image.fmax) continue
        const y = M.t + ph - (v / image.fmax) * ph
        ctx.fillText(formatTick(v), M.l - 5, y)
      }
      ctx.textAlign = 'center'
      ctx.textBaseline = 'top'
      for (const v of niceTicks(0, image.durationSec, 6)) {
        if (v > image.durationSec) continue
        const x = M.l + (v / image.durationSec) * pw
        ctx.fillText(formatTick(v), x, M.t + ph + 4)
      }
      ctx.textBaseline = 'bottom'
      ctx.fillText('time (s)', M.l + pw / 2, h - 1)
      ctx.save()
      ctx.translate(11, M.t + ph / 2)
      ctx.rotate(-Math.PI / 2)
      ctx.textBaseline = 'middle'
      ctx.fillText('Hz', 0, 0)
      ctx.restore()

      ctx.fillStyle = labelColor
      ctx.font = 'bold 12px system-ui, sans-serif'
      ctx.textAlign = 'left'
      ctx.textBaseline = 'top'
      ctx.fillText(label, M.l + 6, M.t + 5)
    },
    [raster, image, label, labelColor, highlight],
  )
  const ref = useCanvas(draw)
  return <canvas ref={ref} className={`plot ${className ?? 'image-plot'}`} />
}