import { useLayoutEffect, useMemo, useState, type ReactNode } from 'react'
import {
  LegendContext,
  createLegendStore,
  getLegendStore,
  hasLegendStore,
  registerLegendStore,
  setLegendStoreEntry,
  unregisterLegendStore,
} from '../stores'
import type {
  LegendGroupInput,
  LegendLayerInput,
  LegendStoreApi,
} from '../stores'
import { LegendConfigContext, type LegendConfig } from './legend-config-context'
import { DEFAULT_LEGEND_LABELS, type LegendLabels } from './labels'

export interface LegendProviderProps {
  /** Unique id for this legend panel — keys its store. */
  id: string
  /**
   * Layer definitions. Interactive fields (visible/opacity/…) are seeded.
   * Pass stable references — `layers`, `groups`, and each layer's `variables`
   * should be memoized or module constants (same contract as widgets-v2
   * `data`); the store bails out of unchanged syncs by reference.
   */
  layers: LegendLayerInput[]
  /** Optional group definitions for grouping layers. */
  groups?: LegendGroupInput[]
  /** Override any subset of the legend's user-facing strings. */
  labels?: Partial<LegendLabels>
  /**
   * When `true` (the default), the per-panel store survives unmount so
   * visibility/opacity/order are preserved across remounts — the common case
   * for a legend that toggles in and out of view. Consumers must call
   * `deleteLegendStore(id)` when the legend is permanently removed. Set to
   * `false` to tear the store down on unmount.
   */
  keepAlive?: boolean
  children: ReactNode
}

/**
 * Legend panel shell. Creates a per-panel Zustand store keyed by `id`, syncs
 * the consumer's `layers`/`groups` definitions onto it (preserving interactive
 * state), and exposes the `id` via `LegendContext` so descendants resolve their
 * store with `useLegendId()`. Labels are surfaced via `LegendConfigContext`.
 *
 * Modeled on the widgets-v2 `Widget.Provider`.
 *
 * @experimental This API is new and may change in a future release.
 */
export function LegendProvider({
  id,
  layers,
  groups,
  labels,
  keepAlive = true,
  children,
}: LegendProviderProps) {
  // Lazy init — reuse an existing keepAlive store, otherwise create one and
  // synchronously place it in the registry so descendants can call
  // useLegendStore(id, ...) on first render.
  const [store] = useState<LegendStoreApi>(() => {
    if (hasLegendStore(id)) return getLegendStore(id)
    const created = createLegendStore(id, { layers, groups })
    setLegendStoreEntry(id, created)
    return created
  })

  useLayoutEffect(() => {
    registerLegendStore(id, store)
    return () => {
      unregisterLegendStore(id, { keepAlive })
    }
  }, [id, store, keepAlive])

  // Sync definitions in. The store's `_sync` preserves interactive state of
  // known layers/groups and bails out (no state change) when the merged result
  // is reference-identical, so stable props are a no-op.
  useLayoutEffect(() => {
    store.getState()._sync(layers, groups ?? [])
  }, [store, layers, groups])

  const config = useMemo<LegendConfig>(
    () => ({
      labels: labels
        ? { ...DEFAULT_LEGEND_LABELS, ...labels }
        : DEFAULT_LEGEND_LABELS,
    }),
    [labels],
  )

  return (
    <LegendContext.Provider value={id}>
      <LegendConfigContext.Provider value={config}>
        {children}
      </LegendConfigContext.Provider>
    </LegendContext.Provider>
  )
}