import {
  useCallback,
  useEffect,
  useMemo,
  useRef,
  type Dispatch,
  type SetStateAction,
} from 'react'
import { useStore } from 'zustand'
import { getWidgetStore } from './widget-store-registry'
import { removeTransform, upsertTransform } from './transforms'
import type {
  Transform,
  TransformDescriptor,
  TransformFn,
  TransformPair,
  WidgetState,
} from './types'

export interface UseTransformOptions {
  initialEnabled?: boolean
}

export interface TransformControls {
  enabled: boolean
  setEnabled: Dispatch<SetStateAction<boolean>>
  toggle: () => void
}

type RegistrationKey = 'dataTransforms' | 'configTransforms'

interface Registration {
  key: RegistrationKey
  id: string
}

export function useTransform(
  id: string,
  pairs: readonly TransformPair[],
  options?: UseTransformOptions,
): TransformControls {
  if (pairs.length === 0) {
    throw new Error('[widgets-v2] useTransform requires at least one pair.')
  }

  const store = getWidgetStore(id)
  const primaryId = pairs[0]!.descriptor.id
  const initialEnabled = options?.initialEnabled ?? true

  // Read `enabled` reactively from the store. Multiple mounts of the same
  // descriptor (e.g., Toolbox renders ZoomToggle both in the inline preview
  // and inside its overflow Paper) stay in sync because they share the
  // `transformStates[primaryId].enabled` slot.
  const persisted = useStore(
    store,
    (s) => s.transformStates[primaryId]?.enabled,
  )
  const enabled = typeof persisted === 'boolean' ? persisted : initialEnabled

  // Seed the store's `transformStates[primaryId].enabled` on first mount so
  // subsequent reads (and a rehydration after `keepAlive` remounts) have a
  // canonical value to read from. Idempotent: skip when already set.
  useEffect(() => {
    if (typeof persisted === 'boolean') return
    store.setState((s: WidgetState) => ({
      transformStates: {
        ...s.transformStates,
        [primaryId]: {
          ...s.transformStates[primaryId],
          enabled: initialEnabled,
        },
      },
    }))
  }, [store, primaryId, persisted, initialEnabled])

  const setEnabled = useCallback<Dispatch<SetStateAction<boolean>>>(
    (next) => {
      store.setState((s: WidgetState) => {
        const current = s.transformStates[primaryId]?.enabled ?? initialEnabled
        const value =
          typeof next === 'function'
            ? (next as (v: boolean) => boolean)(current)
            : next
        if (value === current) return {}
        return {
          transformStates: {
            ...s.transformStates,
            [primaryId]: {
              ...s.transformStates[primaryId],
              enabled: value,
            },
          },
        }
      })
    },
    [store, primaryId, initialEnabled],
  )

  // Track the latest registrations so the unmount cleanup knows what to remove.
  // Updated inside the register effect (post-render), never during render.
  const registeredRef = useRef<Registration[]>([])

  // Register / re-register when `pairs` reference or `enabled` changes.
  // upsertTransform is identity-stable: re-registering with an unchanged
  // descriptor + fn returns the same array reference, and the middleware
  // bypasses the pipeline for unchanged refs. Consumers should memoize
  // `pairs` (useSingleTransform does this internally; multi-pair consumers
  // should wrap the pairs array in useMemo).
  useEffect(() => {
    const next: Registration[] = []
    for (const { descriptor, fn } of pairs) {
      const key: RegistrationKey =
        descriptor.type === 'data' ? 'dataTransforms' : 'configTransforms'
      const transform: Transform = { ...descriptor, fn, enabled }
      store.setState(
        (s: WidgetState) =>
          ({
            [key]: upsertTransform(s[key], transform),
          }) as Partial<WidgetState>,
      )
      next.push({ key, id: descriptor.id })
    }
    registeredRef.current = next
  }, [store, pairs, enabled])

  // Unmount cleanup — remove every currently-registered descriptor. Reads from
  // the ref so we always remove the latest registrations, not a stale snapshot.
  useEffect(() => {
    return () => {
      for (const { key, id: descriptorId } of registeredRef.current) {
        store.setState(
          (s: WidgetState) =>
            ({
              [key]: removeTransform(s[key], descriptorId),
            }) as Partial<WidgetState>,
        )
      }
      registeredRef.current = []
    }
  }, [store])

  const toggle = useCallback(() => setEnabled((v) => !v), [setEnabled])

  return { enabled, setEnabled, toggle }
}

export function useSingleTransform(
  id: string,
  descriptor: TransformDescriptor,
  fn: TransformFn,
  options?: UseTransformOptions,
): TransformControls {
  const pairs = useMemo<readonly TransformPair[]>(
    () => [{ descriptor, fn }],
    [descriptor, fn],
  )
  return useTransform(id, pairs, options)
}