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
}

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(args: {
  isGmaps: boolean
  dpr: number
  windowDpr: number
  deckInstance: DeckInstance
  overlayInstance: NonNullable<OverlayRef>
  maplibreMap: MaplibreMap | null
}): Promise<string> {
  const {
    isGmaps,
    dpr,
    windowDpr,
    deckInstance,
    overlayInstance,
    maplibreMap,
  } = args

  const needsDeckResize =
    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<string> => {
    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 canvas.toDataURL()
  }

  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. Google Maps basemaps are captured via `html2canvas` at the requested scale and are unaffected by the maplibre resize path.
 * @returns An object with `takeScreenshot` to trigger capture, `screenshot` as a data URL string, `isLoading`, `error`, and `reset`.
 *
 * @remarks
 * Google Maps basemaps must use raster rendering mode; vector mode is not supported.
 * 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.
 * The capture additionally waits for the basemap's `idle` event and for `deck.needsRedraw()` to return `false` for two consecutive frames before drawing. This avoids 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 })
 * ```
 */
export function useMapScreenshot({
  mapId,
  pixelRatio,
}: {
  mapId: MapComponentProps['id']
  pixelRatio?: number
}) {
  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 isAllLoaded = isLoaded && areLayersLoaded && isIdle

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

  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)
      }

      // Validate Google Maps rendering type
      if (isGmaps) {
        const _overlayInstance = overlayInstance as google.maps.Map
        if ((_overlayInstance.getRenderingType() as string) !== 'RASTER') {
          throw new Error('Only raster type is supported for Google Maps')
        }
      }

      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({
        isGmaps,
        dpr,
        windowDpr,
        deckInstance,
        overlayInstance,
        maplibreMap: isGmaps ? null : (overlayInstance as MapRef).getMap(),
      })
    },
  })

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

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