import type { ReactElement, ReactNode } from 'react';
import { Children, isValidElement, useCallback, useEffect, useLayoutEffect, useRef, useState } from 'react';
import { PresenceContext } from './animate-presence-context';

function getValidChildren(children: ReactNode): ReactElement[] {
  const valid: ReactElement[] = [];
  Children.forEach(children, (child) => {
    if (isValidElement(child)) valid.push(child);
  });
  return valid;
}

function setsEqual<T>(a: ReadonlySet<T>, b: ReadonlySet<T>): boolean {
  if (a.size !== b.size) return false;
  for (const item of a) if (!b.has(item)) return false;
  return true;
}

/** A copy of `set` with every entry in `keys` added. */
function withKeysAdded(set: ReadonlySet<string>, keys: readonly string[]): Set<string> {
  const next = new Set(set);
  for (const key of keys) next.add(key);
  return next;
}

/** A copy of `set` with every entry in `keys` removed. */
function withKeysRemoved(set: ReadonlySet<string>, keys: readonly string[]): Set<string> {
  const next = new Set(set);
  for (const key of keys) next.delete(key);
  return next;
}

type KeyClassification = {
  newExits: string[];
  abandonedHeldKeys: string[];
  reenteredKeys: string[];
  newEntries: string[];
};

/** Partition the current keys against the previous render into the four sets the lifecycle reacts to. */
function classifyKeyChanges(
  prevCurrentKeys: ReadonlySet<string>,
  currentKeys: ReadonlySet<string>,
  exitingKeys: ReadonlySet<string>,
  heldKeys: ReadonlySet<string>,
): KeyClassification {
  // Keys that left the tree since the last render. A held key was never
  // rendered, so it is dropped silently instead of being exited.
  const newExits: string[] = [];
  const abandonedHeldKeys: string[] = [];
  for (const key of prevCurrentKeys) {
    if (!(currentKeys.has(key) || exitingKeys.has(key))) {
      if (heldKeys.has(key)) abandonedHeldKeys.push(key);
      else newExits.push(key);
    }
  }

  // Keys that re-appeared while exiting — cancel their exit.
  const reenteredKeys: string[] = [];
  for (const key of exitingKeys) if (currentKeys.has(key)) reenteredKeys.push(key);

  // Keys that just appeared for the first time.
  const newEntries: string[] = [];
  for (const key of currentKeys) if (!prevCurrentKeys.has(key)) newEntries.push(key);

  return { newExits, abandonedHeldKeys, reenteredKeys, newEntries };
}

/**
 * Prune keys that have fully unmounted from the stable order, then insert new
 * current keys at the position their sibling order implies, so exiting items
 * stay at their original list position rather than being appended at the end.
 */
function maintainKeyOrder(order: readonly string[], activeKeys: ReadonlySet<string>, currentKeys: readonly string[]): string[] {
  const result = order.filter((key) => activeKeys.has(key));
  const inOrder = new Set(result);
  for (const [index, key] of currentKeys.entries()) {
    if (!inOrder.has(key)) {
      if (index === 0) result.unshift(key);
      else {
        // index > 0 is guaranteed by the else branch above; fall back to append if somehow undefined
        const prevKey = currentKeys[index - 1];
        const prevIdx = prevKey === undefined ? -1 : result.indexOf(prevKey);
        result.splice(prevIdx >= 0 ? prevIdx + 1 : result.length, 0, key);
      }
      inOrder.add(key);
    }
  }
  return result;
}

type KeyCallbacks = { register: (childId: string) => () => void; safeToUnmount: (childId: string) => void };

export type AnimatePresenceProps = {
  children?: ReactNode;
  /** Custom data passed to exit variant functions as the first argument. */
  custom?: unknown;
  /**
   * When `false`, children present on the first render will not animate in.
   * @default true
   */
  initial?: boolean;
  /** Called once all exiting children have completed their exit animations. */
  onExitComplete?: () => void;
  /**
   * When `true`, a new child will not animate in until all exiting children
   * have completed their exit animations.
   * @default false
   */
  exitBeforeEnter?: boolean;
};

/**
 * Orchestrates exit animations for its children before they unmount.
 *
 * Wrap a set of Moti components (each with an `exit` prop) inside
 * `<AnimatePresence>`. When a child leaves the tree, it is kept alive until
 * its exit animation completes — then unmounted. This is the Moti equivalent
 * of Framer Motion's `AnimatePresence`.
 *
 * ## Key requirements
 *
 * - Every direct child **must** have a unique, stable `key`.
 * - Every exiting child **must** have an `exit` prop (otherwise it unmounts
 *   immediately).
 * - Children that use `exit` **must** be Moti components (`MotiView`,
 *   `MotiText`, etc.) — they register with the presence system via
 *   `usePresenceContext()`.
 *
 * ```tsx
 * <AnimatePresence>
 *   {isVisible && (
 *     <MotiView
 *       key="modal"
 *       from={{ opacity: 0, scale: 0.9 }}
 *       animate={{ opacity: 1, scale: 1 }}
 *       exit={{ opacity: 0, scale: 0.9 }}
 *       transition={{ type: 'spring' }}
 *     />
 *   )}
 * </AnimatePresence>
 * ```
 *
 * @see {@link usePresenceContext} for the hook consumed by every Moti component.
 */
// biome-ignore lint/complexity/noExcessiveLinesPerFunction: AnimatePresence orchestrates the full animation lifecycle (animate/from/exit/state/presence) — the remaining lines are derived-state flushes and registry wiring that share refs and setters, not separable concerns
export function AnimatePresence({ children, custom, initial = true, onExitComplete, exitBeforeEnter }: AnimatePresenceProps) {
  const validChildren = getValidChildren(children);

  // Build a key → element map for the current render
  const currentMap = new Map<string, ReactElement>();
  for (const child of validChildren) currentMap.set(child.key ?? '', child);

  // Store of elements we may need to re-render while they exit. Pruned in the
  // cleanup effect once a key is neither current nor exiting.
  const allElementsRef = useRef<Map<string, ReactElement>>(new Map());
  for (const [k, v] of currentMap) allElementsRef.current.set(k, v);

  // Keys of children that are in the middle of their exit animation.
  const [exitingKeys, setExitingKeys] = useState<ReadonlySet<string>>(new Set());

  // Previous set of "current" keys — used to detect removals during render.
  const [prevCurrentKeys, setPrevCurrentKeys] = useState<ReadonlySet<string>>(new Set());

  // When exitBeforeEnter=true, keys that want to enter but must wait for all
  // exits to complete before becoming visible.
  const [heldKeys, setHeldKeys] = useState<ReadonlySet<string>>(new Set());

  const isFirstRender = useRef(true);

  // key → (childId → exit-complete). Every Moti descendant of a key registers
  // here via PresenceContext.register; the key unmounts only when every
  // registered child has reported completion.
  const registryRef = useRef(new Map<string, Map<string, boolean>>());

  // Stable per-key register/safeToUnmount closures. Identity MUST NOT change
  // across renders: children re-register when `register` changes identity, and
  // the transient deregister could complete an in-flight exit prematurely.
  const keyCallbacksRef = useRef(new Map<string, KeyCallbacks>());

  // Tracks insertion order of all keys (current + exiting) so exiting items
  // stay at their original list position instead of being appended at the end.
  const keyOrderRef = useRef<string[]>([]);

  // Removes `key` from the exiting set once every registered child has reported
  // completion (or none are registered). Never called during render.
  const finishExitIfDone = useCallback((key: string) => {
    const registered = registryRef.current.get(key);
    if (registered) for (const done of registered.values()) if (!done) return;
    setExitingKeys((prev) => {
      if (!prev.has(key)) return prev;
      const next = new Set(prev);
      next.delete(key);
      return next;
    });
  }, []);

  const getKeyCallbacks = (key: string): KeyCallbacks => {
    let callbacks = keyCallbacksRef.current.get(key);
    if (!callbacks) {
      callbacks = {
        register: (childId) => {
          let registered = registryRef.current.get(key);
          if (!registered) {
            registered = new Map();
            registryRef.current.set(key, registered);
          }
          registered.set(childId, false);
          return () => {
            registered.delete(childId);
            // A child unmounting mid-exit must not block the key forever.
            finishExitIfDone(key);
          };
        },
        safeToUnmount: (childId) => {
          registryRef.current.get(key)?.set(childId, true);
          finishExitIfDone(key);
        },
      };
      keyCallbacksRef.current.set(key, callbacks);
    }
    return callbacks;
  };

  const currentKeySet: ReadonlySet<string> = new Set(currentMap.keys());

  const { newExits, abandonedHeldKeys, reenteredKeys, newEntries } = classifyKeyChanges(
    prevCurrentKeys,
    currentKeySet,
    exitingKeys,
    heldKeys,
  );

  // Reset completion flags for keys starting or canceling an exit, so stale
  // reports from a previous (canceled) exit can't complete the next one early.
  // Ref mutation during render: idempotent, safe under StrictMode re-renders.
  for (const key of [...newExits, ...reenteredKeys]) {
    const registered = registryRef.current.get(key);
    if (registered) for (const childId of registered.keys()) registered.set(childId, false);
  }

  // Flush state updates synchronously inside render (React's safe pattern
  // for derived state — triggers an immediate additional render pass).
  // biome-ignore lint/plugin: intentional derived-state-in-render (React's getDerivedStateFromProps equivalent); guarded by setsEqual so it converges and cannot loop
  if (!setsEqual(prevCurrentKeys, currentKeySet)) setPrevCurrentKeys(currentKeySet);
  if (newExits.length > 0)
    // biome-ignore lint/plugin: intentional derived-state-in-render; only fires when keys actually leave, converging to a stable exiting set
    setExitingKeys((prev) => withKeysAdded(prev, newExits));
  if (reenteredKeys.length > 0)
    // biome-ignore lint/plugin: intentional derived-state-in-render; only fires when an exiting key re-appears in the tree, converges once the key leaves the exiting set
    setExitingKeys((prev) => withKeysRemoved(prev, reenteredKeys));
  if (abandonedHeldKeys.length > 0)
    // biome-ignore lint/plugin: intentional derived-state-in-render; only fires when a held (never-rendered) key leaves the tree, converges once it is dropped from the held set
    setHeldKeys((prev) => withKeysRemoved(prev, abandonedHeldKeys));

  // exitBeforeEnter: hold new entries while exits are in progress. `newExits`
  // must be counted too — on an atomic key swap the exit detected in this very
  // render pass has not been committed to `exitingKeys` yet (Fix B).
  if (exitBeforeEnter && newEntries.length > 0 && (exitingKeys.size > 0 || newExits.length > 0))
    // biome-ignore lint/plugin: intentional derived-state-in-render; guarded so it only fires when new keys arrive while exits are active, converges to stable held set
    setHeldKeys((prev) => withKeysAdded(prev, newEntries));

  // Release held keys once all exits have completed.
  if (exitBeforeEnter && exitingKeys.size === 0 && heldKeys.size > 0)
    // biome-ignore lint/plugin: intentional derived-state-in-render; only fires when exits drain and held keys remain, converges immediately to empty set
    setHeldKeys(new Set());

  const hasActiveExits = exitingKeys.size > 0;

  // Maintain a stable insertion-order for all active keys so exiting items
  // remain at their original position in the list rather than being appended
  // at the end. Ref mutation during render: idempotent under StrictMode.
  // newExits are detected this render but not yet in exitingKeys state (which
  // updates on the next pass). Include them so keyOrderRef doesn't prune a key
  // before it's been added to exitingKeys, which would prevent its exit render.
  const allActive = new Set([...currentKeySet, ...exitingKeys, ...newExits]);
  keyOrderRef.current = maintainKeyOrder(
    keyOrderRef.current,
    allActive,
    validChildren.map((child) => child.key ?? ''),
  );

  // Build the render list from the stable order, respecting exitBeforeEnter.
  const keysToRender: string[] = [];
  for (const key of keyOrderRef.current) {
    if ((currentKeySet.has(key) || exitingKeys.has(key)) && !(exitBeforeEnter && hasActiveExits && heldKeys.has(key)))
      keysToRender.push(key);
  }

  const firstRender = isFirstRender.current;
  isFirstRender.current = false;

  // Fire onExitComplete when the exiting set drains — but only when the drained
  // keys actually finished rather than being canceled by re-entry. useModalRender
  // unmounts its Modal on this callback; firing it after a rapid re-open would
  // tear down a modal that is open again.
  const prevExitingRef = useRef<ReadonlySet<string>>(new Set());
  // Fire onExitComplete with useLayoutEffect so the callback runs synchronously after
  // commit but before paint. State updates (unmounting a Modal, resetting a lifted
  // trigger) land in the same frame as the exit animation completing — no gap where
  // neither the exiting copy nor the original is visible.
  useLayoutEffect(() => {
    const prevExiting = prevExitingRef.current;
    prevExitingRef.current = exitingKeys;
    if (prevExiting.size === 0 || exitingKeys.size > 0) return;
    for (const key of prevExiting) if (currentKeySet.has(key)) return;
    onExitComplete?.();
  });

  // Safety net + cache cleanup. finishExitIfDone completes exiting keys whose
  // subtree contains no registered Moti children (nothing would ever report),
  // and keys that are neither current nor exiting release their caches so
  // allElementsRef doesn't grow forever (TextCascade mints a key per label).
  // biome-ignore lint/plugin: pruning per-key caches and completing childless exits must run after children's effects have registered/reported — inherently a post-commit side effect
  useEffect(() => {
    for (const key of exitingKeys) finishExitIfDone(key);
    for (const key of allElementsRef.current.keys()) {
      if (!(currentKeySet.has(key) || exitingKeys.has(key))) {
        allElementsRef.current.delete(key);
        registryRef.current.delete(key);
        keyCallbacksRef.current.delete(key);
      }
    }
  });

  const wrappedChildren = keysToRender.map((key) => {
    const element = currentMap.get(key) ?? allElementsRef.current.get(key);
    if (!element) return null;

    const isPresent = currentMap.has(key);
    const callbacks = getKeyCallbacks(key);

    return (
      <PresenceContext.Provider
        key={key}
        value={{
          isPresent,
          register: callbacks.register,
          safeToUnmount: isPresent ? null : callbacks.safeToUnmount,
          custom,
          // Pass initial=false when this is the first render and the caller
          // opted out of entry animations with initial={false}.
          initial: firstRender && !initial ? false : undefined,
        }}
      >
        {element}
      </PresenceContext.Provider>
    );
  });

  return <>{wrappedChildren}</>;
}