import { isValidElement } from 'react'
import { useStore, type StoreApi } from 'zustand'
import { createStore } from 'zustand/vanilla'
import { devtools } from 'zustand/middleware'
import { useShallow } from 'zustand/react/shallow'
import type {
  LegendConfigSection,
  LegendGroup,
  LegendGroupInput,
  LegendInit,
  LegendLayer,
  LegendLayerInput,
  LegendState,
  LegendStoreApi,
} from './types'
import { selectLayersByGroup, selectOrderedGroups } from './selectors'

// Per-panel legend store registry. One store per `<Legend.Provider id>`.
// Mirrors the widgets-v2 registry pattern (keyed Map + refcounted mounts),
// without the chart-pipeline machinery.
const legendStores = new Map<string, LegendStoreApi>()
const legendMountCounts = new Map<string, number>()
const warnedDuplicates = new Set<string>()

const isDev = (): boolean => {
  try {
    return Boolean(
      (import.meta as unknown as { env?: { DEV?: boolean } }).env?.DEV,
    )
  } catch {
    return false
  }
}

const clamp01 = (n: number): number => Math.min(1, Math.max(0, n))

// ─── Pure state helpers ─────────────────────────────────────────────────────

/**
 * Visibility↔collapse coupling. Hiding snapshots the current collapsed state
 * and collapses; showing restores the snapshot (manual choices persist) and
 * defaults to expanded. Idempotent: returns `null` when `visible` is already
 * the requested value, so a double-hide can't overwrite the real snapshot.
 */
function visibilityPatch(
  entity: {
    visible: boolean
    collapsed: boolean
    collapsedBeforeHide?: boolean
  },
  visible: boolean,
): {
  visible: boolean
  collapsed: boolean
  collapsedBeforeHide: boolean | undefined
} | null {
  if (entity.visible === visible) return null
  return visible
    ? // Showing: restore the pre-hide collapsed state (default expand).
      {
        visible,
        collapsed: entity.collapsedBeforeHide ?? false,
        collapsedBeforeHide: undefined,
      }
    : // Hiding: snapshot, then collapse.
      {
        visible,
        collapsed: true,
        collapsedBeforeHide: entity.collapsed,
      }
}

/**
 * Manual collapse. While hidden, the snapshot follows the manual choice so
 * the next show keeps it.
 */
function collapsePatch(
  entity: { visible: boolean },
  collapsed: boolean,
): { collapsed: boolean; collapsedBeforeHide?: boolean } {
  return entity.visible
    ? { collapsed }
    : { collapsed, collapsedBeforeHide: collapsed }
}

/**
 * Reorder one entity within its **ordered** bucket: splice it out, insert at
 * the clamped index, and return the dense `order` rewrites (only the entries
 * whose order actually changes). `null` when the position doesn't change or
 * the id isn't in the bucket. Powers `moveLayer` / `moveGroup`.
 */
function movePatch<T extends { id: string; order: number }>(
  bucket: T[],
  id: string,
  toIndex: number,
): Map<string, number> | null {
  const fromIndex = bucket.findIndex((e) => e.id === id)
  if (fromIndex === -1) return null
  const clamped = Math.max(0, Math.min(bucket.length - 1, toIndex))
  if (clamped === fromIndex) return null
  const next = [...bucket]
  const [moved] = next.splice(fromIndex, 1)
  next.splice(clamped, 0, moved!)
  const orders = new Map<string, number>()
  next.forEach((entity, index) => {
    if (entity.order !== index) orders.set(entity.id, index)
  })
  return orders.size > 0 ? orders : null
}

function updateLayer(
  state: LegendState,
  id: string,
  patch: Partial<LegendLayer>,
): Partial<LegendState> {
  const layer = state.layers[id]
  if (!layer) return {}
  return { layers: { ...state.layers, [id]: { ...layer, ...patch } } }
}

function updateGroup(
  state: LegendState,
  id: string,
  patch: Partial<LegendGroup>,
): Partial<LegendState> {
  const group = state.groups[id]
  if (!group) return {}
  return { groups: { ...state.groups, [id]: { ...group, ...patch } } }
}

/**
 * Merge incoming config-select sections while preserving the user's picked
 * `active` for any section id that still exists and whose pick is still among
 * the new `options` (the structure — ids, titles, options, order — always
 * follows the props). Returns the previous array when the merge lands on the
 * same content, so unchanged syncs keep entity identity (bail-out).
 */
function mergeSections(
  prev: LegendConfigSection[] | undefined,
  next: LegendConfigSection[] | undefined,
): LegendConfigSection[] | undefined {
  if (!next || !prev || prev === next) return next
  const merged = next.map((section) => {
    const prevSection = prev.find((p) => p.id === section.id)
    return prevSection &&
      prevSection.active !== section.active &&
      section.options.includes(prevSection.active)
      ? { ...section, active: prevSection.active }
      : section
  })
  const sameAsPrev =
    prev.length === merged.length &&
    merged.every((m, i) => {
      const p = prev[i]!
      return (
        p.id === m.id &&
        p.title === m.title &&
        p.options === m.options &&
        p.active === m.active
      )
    })
  return sameAsPrev ? prev : merged
}

/**
 * Same-type, same-key React elements with shallow-equal props count as
 * equal, so a ReactNode field (e.g. `LegendGroup.icon`) re-created inline on
 * every render (`icon={<MyIcon />}`) doesn't look "changed" on every sync. A
 * different `key` is treated as a real change, same as a different type.
 */
function reactElementEqual(a: unknown, b: unknown): boolean {
  if (!isValidElement(a) || !isValidElement(b)) return false
  if (a.type !== b.type || a.key !== b.key) return false
  return shallowEqualEntity(
    a.props as Record<string, unknown>,
    b.props as Record<string, unknown>,
  )
}

/**
 * Shallow equality over the union of both objects' keys — `Object.is` per
 * value, except React elements, which fall back to {@link reactElementEqual}.
 */
function shallowEqualEntity<T extends object>(a: T, b: T): boolean {
  const keys = new Set([...Object.keys(a), ...Object.keys(b)]) as Set<keyof T>
  for (const key of keys) {
    if (Object.is(a[key], b[key])) continue
    if (!reactElementEqual(a[key], b[key])) return false
  }
  return true
}

/**
 * Reassign `order` to dense indexes `0..n-1` within each sibling bucket,
 * sorted by current order with ties broken by input position. This is the
 * only boundary where external (possibly duplicated) orders enter, so the
 * stored orders stay dense and unique within each bucket.
 */
function normalizeOrders<T extends { order: number }>(
  entries: { merged: T; index: number }[],
  bucketOf: (e: T) => string | undefined,
): void {
  const buckets = new Map<string | undefined, { merged: T; index: number }[]>()
  for (const entry of entries) {
    const key = bucketOf(entry.merged)
    const bucket = buckets.get(key)
    if (bucket) bucket.push(entry)
    else buckets.set(key, [entry])
  }
  for (const bucket of buckets.values()) {
    bucket
      .sort((a, b) => a.merged.order - b.merged.order || a.index - b.index)
      .forEach((entry, i) => {
        entry.merged.order = i
      })
  }
}

/**
 * Seed/merge one layer: static fields follow the input; interactive state
 * (`visible`/`opacity`/`collapsed`/`order`/`collapsedBeforeHide` + section
 * actives) sticks to the store entry when present. Shared by `_sync` (per
 * entry) and the imperative `setLegendLayer`.
 */
function mergeLayerEntity(
  prev: LegendLayer | undefined,
  input: LegendLayerInput,
  fallbackOrder: number,
): LegendLayer {
  return {
    ...input,
    visible: prev?.visible ?? input.visible ?? true,
    opacity: prev?.opacity ?? input.opacity ?? 1,
    collapsed: prev?.collapsed ?? input.collapsed ?? false,
    collapsedBeforeHide: prev?.collapsedBeforeHide,
    sections: mergeSections(prev?.sections, input.sections),
    order: prev?.order ?? input.order ?? fallbackOrder,
  } as LegendLayer
}

/**
 * Group counterpart of {@link mergeLayerEntity}. A group with no `prev`
 * state (first sync) defaults collapsed when none of its member layers is
 * visible yet — `hasVisibleLayer` is ignored once the group exists so a
 * user's manual expand/collapse always sticks.
 */
function mergeGroupEntity(
  prev: LegendGroup | undefined,
  input: LegendGroupInput,
  fallbackOrder: number,
  hasVisibleLayer: boolean,
): LegendGroup {
  return {
    ...input,
    collapsed: prev?.collapsed ?? input.collapsed ?? !hasVisibleLayer,
    collapsedBeforeHide: prev?.collapsedBeforeHide,
    order: prev?.order ?? input.order ?? fallbackOrder,
  } as LegendGroup
}

/**
 * Merge incoming definitions while preserving interactive state of known
 * entities: static fields (name, subtitle, variables, attributes, …) follow
 * the props; `visible`/`opacity`/`collapsed`/`order` stick to what's in the
 * store. New entities are seeded; vanished ones are dropped. Sibling orders
 * are normalized to dense indexes.
 *
 * Unchanged entities keep their previous object identity (`variables`
 * compared by reference), and when nothing changed at all the sync is a
 * no-op (`{}`), so stable consumer props never produce a state update.
 */
function syncState(
  state: LegendState,
  layerInputs: LegendLayerInput[],
  groupInputs: LegendGroupInput[],
): Partial<LegendState> {
  const layerEntries = layerInputs.map((input, index) => ({
    index,
    merged: mergeLayerEntity(state.layers[input.id], input, index),
  }))
  normalizeOrders(layerEntries, (l) => l.groupId)

  const groupsWithVisibleLayer = new Set(
    layerEntries
      .filter(({ merged }) => merged.visible && merged.groupId != null)
      .map(({ merged }) => merged.groupId!),
  )
  const groupEntries = groupInputs.map((input, index) => ({
    index,
    merged: mergeGroupEntity(
      state.groups[input.id],
      input,
      index,
      groupsWithVisibleLayer.has(input.id),
    ),
  }))
  normalizeOrders(groupEntries, () => undefined)

  // Reference preservation: reuse the previous entity object when nothing
  // about it changed, so subscribers see stable identities.
  let layersChanged = layerEntries.length !== Object.keys(state.layers).length
  const layers: Record<string, LegendLayer> = {}
  for (const { merged } of layerEntries) {
    const prev = state.layers[merged.id]
    if (prev && shallowEqualEntity(prev, merged)) {
      layers[merged.id] = prev
    } else {
      layers[merged.id] = merged
      layersChanged = true
    }
  }

  let groupsChanged = groupEntries.length !== Object.keys(state.groups).length
  const groups: Record<string, LegendGroup> = {}
  for (const { merged } of groupEntries) {
    const prev = state.groups[merged.id]
    if (prev && shallowEqualEntity(prev, merged)) {
      groups[merged.id] = prev
    } else {
      groups[merged.id] = merged
      groupsChanged = true
    }
  }

  if (!layersChanged && !groupsChanged) return {}
  return {
    layers: layersChanged ? layers : state.layers,
    groups: groupsChanged ? groups : state.groups,
  }
}

/**
 * Rebuild a layer/group record from a candidate map, normalizing dense order
 * per bucket and preserving the previous object identity for entries that are
 * unchanged. Returns `null` when nothing changed (so callers can bail with
 * `{}`).
 */
function rebuildRecord<
  T extends { id: string; order: number; groupId?: string },
>(
  prev: Record<string, T>,
  next: Record<string, T>,
  bucketOf: (e: T) => string | undefined,
): Record<string, T> | null {
  const entries = Object.values(next).map((e, index) => ({
    index,
    merged: { ...e },
  }))
  normalizeOrders(entries, bucketOf)
  let changed = entries.length !== Object.keys(prev).length
  const out: Record<string, T> = {}
  for (const { merged } of entries) {
    const before = prev[merged.id]
    if (before && shallowEqualEntity(before, merged)) {
      out[merged.id] = before
    } else {
      out[merged.id] = merged
      changed = true
    }
  }
  return changed ? out : null
}

/** Add or update a single layer (imperative `setLegendLayer`). */
function upsertLayerState(
  state: LegendState,
  input: LegendLayerInput,
): Partial<LegendState> {
  const siblingCount = Object.values(state.layers).filter(
    (l) => l.groupId === input.groupId && l.id !== input.id,
  ).length
  const merged = mergeLayerEntity(state.layers[input.id], input, siblingCount)
  const layers = rebuildRecord(
    state.layers,
    { ...state.layers, [input.id]: merged },
    (l) => l.groupId,
  )
  return layers ? { layers } : {}
}

/** Remove a single layer, re-densifying its sibling orders. */
function removeLayerState(
  state: LegendState,
  id: string,
): Partial<LegendState> {
  if (!state.layers[id]) return {}
  const next = { ...state.layers }
  delete next[id]
  const layers = rebuildRecord(state.layers, next, (l) => l.groupId)
  return layers ? { layers } : {}
}

/** Add or update a single group (imperative `setLegendGroup`). */
function upsertGroupState(
  state: LegendState,
  input: LegendGroupInput,
): Partial<LegendState> {
  const count = Object.values(state.groups).filter(
    (g) => g.id !== input.id,
  ).length
  const hasVisibleLayer = Object.values(state.layers).some(
    (l) => l.groupId === input.id && l.visible,
  )
  const merged = mergeGroupEntity(
    state.groups[input.id],
    input,
    count,
    hasVisibleLayer,
  )
  const groups = rebuildRecord(
    state.groups,
    { ...state.groups, [input.id]: merged },
    () => undefined,
  )
  return groups ? { groups } : {}
}

/**
 * Remove a group; its member layers are un-grouped (`groupId` cleared) so they
 * survive as ungrouped rows. Re-densifies both group and layer orders.
 */
function removeGroupState(
  state: LegendState,
  id: string,
): Partial<LegendState> {
  if (!state.groups[id]) return {}
  const nextGroups = { ...state.groups }
  delete nextGroups[id]
  const groups = rebuildRecord(state.groups, nextGroups, () => undefined)

  const members = Object.values(state.layers).filter((l) => l.groupId === id)
  if (members.length === 0) return groups ? { groups } : {}

  const nextLayers = { ...state.layers }
  for (const member of members) {
    nextLayers[member.id] = { ...member, groupId: undefined }
  }
  const layers = rebuildRecord(state.layers, nextLayers, (l) => l.groupId)
  return {
    ...(groups ? { groups } : {}),
    ...(layers ? { layers } : {}),
  }
}

// ─── Store factory ────────────────────────────────────────────────────────────

/**
 * @internal — used by Provider; not part of the public API.
 *
 * @experimental This API is new and may change in a future release.
 */
export function createLegendStore(
  id: string,
  init: LegendInit = {},
): LegendStoreApi {
  // Seed through the same merge/normalize path as `_sync`, so a later sync
  // with identical input is a guaranteed no-op (stable identities).
  const seeded = syncState(
    { layers: {}, groups: {} } as LegendState,
    init.layers ?? [],
    init.groups ?? [],
  )

  return createStore<LegendState>()(
    devtools(
      (set) => ({
        layers: seeded.layers ?? {},
        groups: seeded.groups ?? {},

        setVisibility: (lid, visible) =>
          set(
            (s) => {
              const layer = s.layers[lid]
              const patch = layer && visibilityPatch(layer, visible)
              return patch ? updateLayer(s, lid, patch) : {}
            },
            false,
            'setVisibility',
          ),
        setOpacity: (lid, opacity) =>
          set(
            (s) => updateLayer(s, lid, { opacity: clamp01(opacity) }),
            false,
            'setOpacity',
          ),
        setCollapsed: (lid, collapsed) =>
          set(
            (s) => {
              const layer = s.layers[lid]
              return layer
                ? updateLayer(s, lid, collapsePatch(layer, collapsed))
                : {}
            },
            false,
            'setCollapsed',
          ),
        setSectionValue: (lid, sectionId, value) =>
          set(
            (s) => {
              const layer = s.layers[lid]
              const section = layer?.sections?.find((x) => x.id === sectionId)
              if (
                !layer?.sections ||
                !section ||
                section.active === value ||
                !section.options.includes(value)
              ) {
                return {}
              }
              const sections = layer.sections.map((x) =>
                x.id === sectionId ? { ...x, active: value } : x,
              )
              return updateLayer(s, lid, { sections })
            },
            false,
            'setSectionValue',
          ),

        moveLayer: (lid, toIndex) =>
          set(
            (s) => {
              const layer = s.layers[lid]
              if (!layer) return {}
              const bucket = selectLayersByGroup(s, layer.groupId)
              const orders = movePatch(bucket, lid, toIndex)
              if (!orders) return {}
              const layers = { ...s.layers }
              for (const [id, order] of orders) {
                layers[id] = { ...layers[id]!, order }
              }
              return { layers }
            },
            false,
            'moveLayer',
          ),

        setGroupVisibility: (gid, visible) =>
          set(
            (s) => {
              const group = s.groups[gid]
              if (!group) return {}
              const members = Object.values(s.layers).filter(
                (l) => l.groupId === gid,
              )
              const anyVisible = members.some((l) => l.visible)
              // Idempotence on the derived state: hide needs a visible member,
              // show needs a fully-hidden group (mixed acts in both directions
              // except show-from-mixed still shows the hidden remainder).
              if (members.length === 0 || (!visible && !anyVisible)) return {}

              // Cascade: each member runs the standard visibility↔collapse
              // coupling, so per-layer collapse snapshots survive a group
              // hide/show; members already at the target are untouched.
              const layers = { ...s.layers }
              let layersChanged = false
              for (const member of members) {
                const patch = visibilityPatch(member, visible)
                if (patch) {
                  layers[member.id] = { ...member, ...patch }
                  layersChanged = true
                }
              }
              if (!layersChanged) return {}

              // Group-level collapse coupling, only on real derived
              // transitions: any-visible→hidden snapshots + collapses;
              // hidden→visible restores the snapshot.
              let groupPatch: Partial<LegendGroup> | null = null
              if (!visible && anyVisible) {
                groupPatch = {
                  collapsed: true,
                  collapsedBeforeHide: group.collapsed,
                }
              } else if (visible && !anyVisible) {
                groupPatch = {
                  collapsed: group.collapsedBeforeHide ?? false,
                  collapsedBeforeHide: undefined,
                }
              }
              return {
                layers,
                ...(groupPatch
                  ? {
                      groups: {
                        ...s.groups,
                        [gid]: { ...group, ...groupPatch },
                      },
                    }
                  : {}),
              }
            },
            false,
            'setGroupVisibility',
          ),
        setGroupCollapsed: (gid, collapsed) =>
          set(
            (s) => {
              const group = s.groups[gid]
              if (!group) return {}
              // Group visibility is derived from members; while the group is
              // hidden the snapshot follows the manual choice (layer parity).
              const members = Object.values(s.layers).filter(
                (l) => l.groupId === gid,
              )
              const hidden =
                members.length > 0 && members.every((l) => !l.visible)
              return updateGroup(
                s,
                gid,
                hidden
                  ? { collapsed, collapsedBeforeHide: collapsed }
                  : { collapsed },
              )
            },
            false,
            'setGroupCollapsed',
          ),

        moveGroup: (gid, toIndex) =>
          set(
            (s) => {
              if (!s.groups[gid]) return {}
              const orders = movePatch(selectOrderedGroups(s), gid, toIndex)
              if (!orders) return {}
              const groups = { ...s.groups }
              for (const [id, order] of orders) {
                groups[id] = { ...groups[id]!, order }
              }
              return { groups }
            },
            false,
            'moveGroup',
          ),

        setLegendLayer: (input) =>
          set((s) => upsertLayerState(s, input), false, 'setLegendLayer'),
        removeLegendLayer: (lid) =>
          set((s) => removeLayerState(s, lid), false, 'removeLegendLayer'),
        setLegendGroup: (input) =>
          set((s) => upsertGroupState(s, input), false, 'setLegendGroup'),
        removeLegendGroup: (gid) =>
          set((s) => removeGroupState(s, gid), false, 'removeLegendGroup'),

        _sync: (layerInputs, groupInputs) =>
          set((s) => syncState(s, layerInputs, groupInputs), false, '_sync'),
      }),
      { name: `legend-${id}`, enabled: isDev() },
    ),
  )
}

// ─── Registry ───────────────────────────────────────────────────────────────

/** @experimental This API is new and may change in a future release. */
export function getLegendStore(id: string): LegendStoreApi {
  const store = legendStores.get(id)
  if (!store) throw new Error(`[legend] Legend store "${id}" not found.`)
  return store
}

/** @experimental This API is new and may change in a future release. */
export function hasLegendStore(id: string): boolean {
  return legendStores.has(id)
}

/** @experimental This API is new and may change in a future release. */
export function deleteLegendStore(id: string): void {
  legendStores.delete(id)
}

/** @experimental This API is new and may change in a future release. */
export function clearAllLegendStores(): void {
  legendStores.clear()
  legendMountCounts.clear()
  warnedDuplicates.clear()
}

/**
 * @internal — used by Provider's lazy init so children can resolve the store first render.
 *
 * @experimental This API is new and may change in a future release.
 */
export function setLegendStoreEntry(id: string, store: LegendStoreApi): void {
  legendStores.set(id, store)
}

/**
 * @internal — registers a Provider mount and warns on duplicate ids in dev.
 *
 * @experimental This API is new and may change in a future release.
 */
export function registerLegendStore(id: string, store: LegendStoreApi): void {
  legendStores.set(id, store)
  const count = (legendMountCounts.get(id) ?? 0) + 1
  legendMountCounts.set(id, count)
  if (count > 1 && isDev() && !warnedDuplicates.has(id)) {
    // eslint-disable-next-line no-console
    console.warn(
      `[legend] Duplicate <Legend.Provider id="${id}"> detected. ` +
        `Multiple providers sharing an id will race on prop sync. ` +
        `Use unique ids per legend panel.`,
    )
    warnedDuplicates.add(id)
  }
}

/**
 * @internal — Provider unmount counterpart; removes the entry on last unmount.
 *
 * @experimental This API is new and may change in a future release.
 */
export function unregisterLegendStore(
  id: string,
  options: { keepAlive: boolean },
): void {
  const count = (legendMountCounts.get(id) ?? 1) - 1
  if (count <= 0) {
    legendMountCounts.delete(id)
    if (!options.keepAlive) legendStores.delete(id)
  } else {
    legendMountCounts.set(id, count)
  }
}

// ─── Hooks ────────────────────────────────────────────────────────────────────

/** @experimental This API is new and may change in a future release. */
export function useLegendStore<T>(
  id: string,
  selector: (state: LegendState) => T,
): T {
  const store = getLegendStore(id) as unknown as StoreApi<LegendState>
  return useStore(store, selector)
}

/** @experimental This API is new and may change in a future release. */
export function useLegendShallow<T>(
  id: string,
  selector: (state: LegendState) => T,
): T {
  const store = getLegendStore(id) as unknown as StoreApi<LegendState>
  return useStore(store, useShallow(selector))
}