import { useCallback } from 'react'
import { useMutation } from '@tanstack/react-query'
import { useMapStore } from './providers/map/provider'
import type { InstanceRef, MapComponentProps, OverlayRef } from './types'
import type { MapRef } from 'react-map-gl/maplibre'
import { MAPS_BASEMAP, MapsTypes } from './const'
import html2canvas from 'html2canvas'

// The deck instance returned by `getDeckInstance` is `Deck<ViewsT> | undefined`
// (see InstanceRef in types.ts). Inside the capture flow it's already been
// narrowed to non-null — alias here so the helper takes the same concrete type
// the rest of the library passes around.
type DeckInstance = NonNullable<InstanceRef['deck']>
// `MapRef.getMap()` returns the underlying maplibre Map instance; we only ever
// call resize/repaint methods on it. Use the real return type so we share the
// store's typing surface rather than redeclaring a Like shim.
type MaplibreMap = ReturnType<MapRef['getMap']>

const MAP_NOT_READY_ERROR = 'Map instances are not ready'
const MAX_RETRIES = 10
const RETRY_DELAY_MS = 500
const DECK_IDLE_TIMEOUT_MS = 2000
const DECK_IDLE_STABLE_FRAMES = 2

interface DeckIdlePollable {
  needsRedraw: (opts?: { clearRedrawFlags: boolean }) => string | false
}

/** Extra capture info handed to `processCapture` alongside the canvas. An object so future fields can be added without breaking callers. */
export interface MapScreenshotCaptureContext {
  /** Device-pixel ratio the canvas was rendered at — use to scale CSS-pixel offsets (e.g. from `getBoundingClientRect()`) into canvas space. On the vector-Google-Maps path this is `window.devicePixelRatio`, since Google's interleaved canvas is captured 1:1 at native device resolution. */
  dpr: number
}

/**
 * Serialization extension point: receives the final composited canvas (basemap
 * + deck.gl layers) right before it would be read via `toDataURL()`, on every
 * capture path (MapLibre/CARTO, raster and vector Google Maps). Draw onto it
 * (logos, attribution) and/or return the final result — a data-URL string or a
 * `Blob` — to replace the default `canvas.toDataURL()` serialization. Return
 * nothing to keep the default. If you return a `Blob`, always return one: a
 * mixed `Blob | void` callback falls back to a string at runtime while being
 * typed as `Blob`.
 */
export type MapScreenshotProcessCapture<
  TResult extends string | Blob = string,
> = (
  canvas: HTMLCanvasElement,
  context: MapScreenshotCaptureContext,
) => TResult | void | Promise<TResult | void>

function waitForDeckIdle(deck: DeckIdlePollable): Promise<void> {
  // Fast path: deck is already idle, no need to wait any frames. The screenshot
  // pipeline calls `deck.redraw('screenshot')` before reading the canvas, so a
  // single synchronous "false" from needsRedraw is sufficient — any redraw that
  // gets queued between now and then will be picked up by that forced redraw.
  if (deck.needsRedraw({ clearRedrawFlags: false }) === false) {
    return Promise.resolve()
  }

  return new Promise((resolve) => {
    let stable = 0
    const start = performance.now()
    const tick = () => {
      if (performance.now() - start >= DECK_IDLE_TIMEOUT_MS) {
        resolve()
        return
      }
      if (deck.needsRedraw({ clearRedrawFlags: false }) === false) {
        stable += 1
        if (stable >= DECK_IDLE_STABLE_FRAMES) {
          resolve()
          return
        }
      } else {
        stable = 0
      }
      requestAnimationFrame(tick)
    }
    requestAnimationFrame(tick)
  })
}

/**
 * Module-level capture flow. Lives outside the hook on purpose: it contains
 * `try/finally` + an imperatively mutated cleanups array, which React Compiler
 * declines to memoize inside a hook body. Pulling it out keeps the hook itself
 * compiler-friendly while leaving the resize/restore semantics intact.
 */
async function captureMapScreenshot<
  TResult extends string | Blob = string,
>(args: {
  gmapsRenderingType: 'RASTER' | 'VECTOR' | null
  dpr: number
  windowDpr: number
  deckInstance: DeckInstance
  overlayInstance: NonNullable<OverlayRef>
  maplibreMap: MaplibreMap | null
  processCapture?: MapScreenshotProcessCapture<TResult>
}): Promise<TResult> {
  const {
    gmapsRenderingType,
    dpr,
    windowDpr,
    deckInstance,
    overlayInstance,
    maplibreMap,
    processCapture,
  } = args

  // Shared serialization step for every capture path. Errors from
  // `processCapture` propagate to the mutation (surfacing consumer bugs); the
  // DPR-restore `finally` below still runs regardless. The cast covers the
  // documented Blob+void mismatch — without `processCapture`, TResult = string.
  const finalize = async (
    c: HTMLCanvasElement,
    effectiveDpr: number,
  ): Promise<TResult> => {
    const result = await processCapture?.(c, { dpr: effectiveDpr })
    return result ?? (c.toDataURL() as TResult)
  }

  const isGmaps = gmapsRenderingType !== null
  const isVector = gmapsRenderingType === 'VECTOR'

  // Vector gmaps captures Google's own canvas 1:1 at its native device
  // resolution (see drawMapAndDeckToCanvas), so the DPR resize dance can't and
  // shouldn't apply — deck doesn't own that canvas.
  const needsDeckResize =
    !isVector &&
    Math.abs(dpr - windowDpr) > 1e-3 &&
    typeof deckInstance.setProps === 'function'
  const needsMaplibreResize =
    Math.abs(dpr - windowDpr) > 1e-3 &&
    maplibreMap !== null &&
    typeof maplibreMap.setPixelRatio === 'function'

  const originalDeckUseDevicePixels =
    deckInstance.props?.useDevicePixels ?? true
  const originalMapPixelRatio = needsMaplibreResize
    ? maplibreMap.getPixelRatio()
    : null

  const cssWidth = deckInstance.width ?? 0
  const cssHeight = deckInstance.height ?? 0
  const canvas = document.createElement('canvas')
  canvas.width = Math.round(cssWidth * dpr)
  canvas.height = Math.round(cssHeight * dpr)
  const context = canvas.getContext('2d')

  const drawMapAndDeckToCanvas = async (): Promise<TResult> => {
    if (isVector) {
      // Interleaved vector gmaps: deck renders into Google's shared WebGL2
      // canvas, which already holds basemap + layers. Capture that single canvas
      // 1:1 at Google's native device resolution — no html2canvas, no compositing
      // (deck.width/height are null here, so we size from the source canvas).
      deckInstance.redraw('screenshot')
      const src = deckInstance.getCanvas()!
      const out = document.createElement('canvas')
      out.width = src.width
      out.height = src.height
      out.getContext('2d')?.drawImage(src, 0, 0)
      return finalize(out, windowDpr)
    }

    if (isGmaps) {
      const _overlayInstance = overlayInstance as google.maps.Map
      const htmlCanvasResult = await html2canvas(_overlayInstance.getDiv(), {
        useCORS: true,
        scale: dpr,
      })
      context?.drawImage(htmlCanvasResult, 0, 0, canvas.width, canvas.height)
    } else {
      const _overlayInstance = overlayInstance as MapRef
      _overlayInstance.redraw()
      context?.drawImage(
        _overlayInstance.getMap().getCanvas(),
        0,
        0,
        canvas.width,
        canvas.height,
      )
    }

    deckInstance.redraw('screenshot')
    context?.drawImage(
      deckInstance.getCanvas()!,
      0,
      0,
      canvas.width,
      canvas.height,
    )

    return finalize(canvas, dpr)
  }

  if (!needsDeckResize && !needsMaplibreResize) {
    return drawMapAndDeckToCanvas()
  }

  // Sequential setup: maplibre first (slow link — waits for tile reloads),
  // then deck (synchronous-ish). Each step pushes its restore into `cleanups`
  // only after completing successfully so a partial failure (throw between
  // steps) only reverses what was done.
  const cleanups: (() => void)[] = []
  try {
    if (needsMaplibreResize) {
      maplibreMap.setPixelRatio(dpr)
      // `setPixelRatio` calls `resize()` internally, which should queue a
      // render → `idle`. `triggerRepaint` belt-and-suspenders: if a maplibre
      // version short-circuits the queued render, the dirty flag still
      // guarantees a frame goes through.
      maplibreMap.triggerRepaint()
      await new Promise<void>((resolve) => {
        void maplibreMap.once('idle', () => resolve())
      })
      cleanups.push(() => {
        maplibreMap.setPixelRatio(originalMapPixelRatio!)
        maplibreMap.triggerRepaint()
      })
    }

    if (needsDeckResize) {
      deckInstance.setProps({ useDevicePixels: dpr })
      // `setProps` only flags the canvas-surface for a drawing-buffer resize;
      // the actual `canvas.width / canvas.height` assignment runs inside
      // `_resizeDrawingBufferIfNeeded`, which is only called from
      // `getCurrentFramebuffer` — i.e. during a render. Force a draw now so
      // the canvas is at the requested DPR before we read its pixels.
      deckInstance.redraw('screenshot-resize')
      await waitForDeckIdle(deckInstance as unknown as DeckIdlePollable)
      cleanups.push(() => {
        deckInstance.setProps({
          useDevicePixels: originalDeckUseDevicePixels,
        })
      })
    }

    return await drawMapAndDeckToCanvas()
  } finally {
    // Fire-and-forget restore: callers get the screenshot as soon as it's
    // captured; maplibre / deck transition back to native DPR in the
    // background. Iterate in reverse so the teardown mirrors setup (LIFO):
    // deck restores before maplibre, minimizing the transient DPR mismatch
    // window during the brief restore-time flicker.
    for (let i = cleanups.length - 1; i >= 0; i--) {
      try {
        cleanups[i]!()
      } catch {
        /* non-fatal — capture already succeeded */
      }
    }
  }
}

/**
 * Captures screenshots of a map instance, supporting both CARTO and Google Maps basemaps.
 *
 * @param mapId - Identifier of the map to capture.
 * @param pixelRatio - Target device-pixel ratio for the capture (both the maplibre basemap and the deck.gl overlay). Defaults to `Math.max(2, window.devicePixelRatio)` — at least 2× so the map is sharp on Linux/Windows displays where `window.devicePixelRatio` is typically 1, but never below the display's native DPR (e.g. 3 on a high-DPI laptop). When the resolved value already matches `window.devicePixelRatio`, no resize happens. Raster Google Maps basemaps are captured via `html2canvas` at the requested scale and are unaffected by the maplibre resize path. **For vector Google Maps `pixelRatio` is a no-op** — the basemap and overlay share Google's WebGL2 canvas, which is captured 1:1 at Google's native device resolution.
 * @param processCapture - Optional serialization extension point, invoked with `(canvas, { dpr })` — the final composited canvas — right before the default `toDataURL()` read, on **every** capture path (MapLibre/CARTO, raster and vector Google Maps). Draw onto the canvas (e.g. a logo, attribution text) and/or return the final result (a data-URL string or a `Blob`) to replace the default serialization; return nothing to keep it. Exceptions from the callback surface through the hook's `error`. On the vector-Google-Maps path `dpr` is `window.devicePixelRatio` (Google's interleaved canvas is captured 1:1 at native device resolution).
 * @returns An object with `takeScreenshot` to trigger capture, `screenshot` as a data URL string (or whatever `processCapture` returns), `isLoading`, `error`, and `reset`.
 *
 * @remarks
 * Both raster and vector Google Maps rendering modes are supported. Raster captures the DOM basemap tiles via `html2canvas` and composites the deck overlay on top; vector reads Google's interleaved WebGL2 canvas directly (basemap + layers in one canvas). A Google map still reporting `UNINITIALIZED` rendering type is treated as not-ready and retried.
 * If map instances are not ready when `takeScreenshot` is called, the capture is automatically retried until they become available.
 * When resizing for capture, maplibre is bumped first (`setPixelRatio` + `triggerRepaint` + `await once('idle')`) so basemap tiles re-render at the new DPR; then deck is bumped via `setProps({ useDevicePixels })` followed by a forced redraw (luma's `_resizeDrawingBufferIfNeeded` only runs from `getCurrentFramebuffer` — i.e. inside a render pass — so `setProps` alone marks the resize but doesn't apply it). After capture, both are restored fire-and-forget so the screenshot promise resolves as soon as the data URL is ready.
 * Readiness gating differs by basemap. MapLibre/CARTO capture waits for the basemap's `idle` flag (`areMapsIdle`) in addition to `isLoaded` and layer load. Google Maps does **not** wait on `areMapsIdle` — it never reliably settles for gmaps because the deck<->Google camera sync fires a trailing `setViewState` after Google's `idle` event and Google doesn't re-fire idle on a static map; gmaps readiness is instead gated by `isLoaded` (tiles), layer load, and `waitForDeckIdle`. In all cases the capture also waits for `deck.needsRedraw()` to return `false` for two consecutive frames before drawing, avoiding partially-painted frames immediately after viewstate changes (e.g. zoom) where tile streaming or layer interpolation can leave the canvas in a transient cut state.
 *
 * @example
 * ```tsx
 * // Default — at least 2× on every OS, native DPR on higher-density displays.
 * const { takeScreenshot, screenshot } = useMapScreenshot({ mapId: 'main-map' })
 *
 * // Force a specific ratio.
 * const { takeScreenshot, screenshot } = useMapScreenshot({ mapId: 'main-map', pixelRatio: 3 })
 *
 * // Draw an overlay and/or take over serialization.
 * const { takeScreenshot, screenshot } = useMapScreenshot({
 *   mapId: 'main-map',
 *   processCapture: (canvas, { dpr }) => {
 *     const ctx = canvas.getContext('2d')!
 *     ctx.drawImage(logo, 16 * dpr, 16 * dpr)
 *     return canvas.toDataURL('image/jpeg', 0.9) // optional — omit to keep the default PNG
 *   },
 * })
 * ```
 */
export function useMapScreenshot<TResult extends string | Blob = string>({
  mapId,
  pixelRatio,
  processCapture,
}: {
  mapId: MapComponentProps['id']
  pixelRatio?: number
  processCapture?: MapScreenshotProcessCapture<TResult>
}) {
  const basemap = useMapStore((state) => state.maps[mapId]?.basemap)!
  const deckInstanceFn = useMapStore((state) => state.getDeckInstance)
  const overlayInstanceFn = useMapStore((state) => state.getOverlayInstance)

  const areLayersLoaded = useMapStore((state) => state.areLayersLoaded[mapId])!
  const isLoaded = useMapStore((state) => state.maps[mapId]?.isLoaded)!
  const isIdle = useMapStore((state) => state.areMapsIdle[mapId])!

  const { type } = MAPS_BASEMAP[basemap]
  const isGmaps = type === MapsTypes.googleMaps

  // `areMapsIdle` is unreliable for Google Maps: the deck<->gmaps camera sync
  // fires a trailing `setViewState` after Google's `idle` event, which resets
  // the idle flag to false and it never recovers (Google doesn't re-fire idle
  // on a static map). Tile load (`isLoaded`), layer load, and the per-capture
  // `waitForDeckIdle` already gate readiness, so skip the idle flag for gmaps.
  const isAllLoaded = isLoaded && areLayersLoaded && (isGmaps || isIdle)

  const mutation = useMutation({
    retry: (failureCount, error) =>
      error.message === MAP_NOT_READY_ERROR && failureCount < MAX_RETRIES,
    retryDelay: RETRY_DELAY_MS,
    mutationFn: async () => {
      const deckInstance = deckInstanceFn(mapId)
      const overlayInstance = overlayInstanceFn(mapId)

      if (!isAllLoaded || !deckInstance || !overlayInstance) {
        throw new Error(MAP_NOT_READY_ERROR)
      }

      // Resolve Google Maps rendering type. Both RASTER and VECTOR are
      // supported; UNINITIALIZED means Google hasn't finished setting up the
      // render pipeline yet, so treat it as not-ready and let the retry wait.
      let gmapsRenderingType: 'RASTER' | 'VECTOR' | null = null
      if (isGmaps) {
        const rt = (
          overlayInstance as google.maps.Map
        ).getRenderingType() as string
        if (rt === 'UNINITIALIZED') {
          throw new Error(MAP_NOT_READY_ERROR)
        }
        gmapsRenderingType = rt as 'RASTER' | 'VECTOR'
      }

      await waitForDeckIdle(deckInstance as unknown as DeckIdlePollable)

      const windowDpr = window.devicePixelRatio ?? 1
      const dpr =
        pixelRatio !== undefined
          ? Math.max(pixelRatio, 1)
          : Math.max(2, windowDpr)

      return captureMapScreenshot({
        gmapsRenderingType,
        dpr,
        windowDpr,
        deckInstance,
        overlayInstance,
        maplibreMap: isGmaps ? null : (overlayInstance as MapRef).getMap(),
        processCapture,
      })
    },
  })

  const takeScreenshot = useCallback(() => {
    mutation.reset()
    mutation.mutate()
  }, [mutation])

  return {
    takeScreenshot,
    screenshot: mutation.data,
    isLoading: mutation.isPending,
    error: mutation.error ?? undefined,
    reset: mutation.reset,
  }
}
