'use client';

/**
 * The React binding for reactive atoms — the client half of the reactivity
 * bridge. The atoms themselves (`atom`, `derive`) are server-safe and live in
 * `../reactivity-core`; this module adds the hooks that read them in a component
 * (`observe`, `<Observe>`, `useAtom*`) and the in-render reader the interpreter
 * uses for a yielded atom. Owning a signal in a service needs no React; only
 * *reading one in the render pass* does — which is why this half is `'use client'`.
 *
 *   const n = useAtomValue(count);                  // read one atom (re-renders on change)
 *   const [n, setN] = useAtom(count);               // read + write, useState-shaped
 *   <Observe>{($) => <b>{$(count)}</b>}</Observe>   // an inline reactive expression in JSX
 */
import { useCallback, useDeferredValue, useRef, useSyncExternalStore, type ReactNode } from 'react';
import type { Atom, ReadableAtom } from '#domain/protocol.ts';
import type { Reader } from '#application/ports.ts';
import {
  computation,
  isDeferred,
  type Computation,
  type Read,
} from '#infrastructure/reactivity-core.ts';

/**
 * Subscribe a component to an inline derived view over one or more atoms. Re-renders
 * precisely when a read atom changes. This is the escape hatch for an ad-hoc reactive
 * expression in a component; prefer reading a single atom with {@link useAtomValue}
 * (or `yield*` in a REC), and deriving shared values with `atom.derive` / `derive` in
 * a service. `<Observe>` is the same thing, inline in JSX.
 *
 * ```tsx
 * const doubled = observe(($) => $(count) * 2);
 * ```
 */
export const observe = <A,>(selector: (read: Read) => A): A => {
  const selectorRef = useRef(selector);
  selectorRef.current = selector;
  const compRef = useRef<Computation<A> | null>(null);
  if (compRef.current === null) {
    compRef.current = computation((read) => selectorRef.current(read));
  }
  const comp = compRef.current;
  const subscribe = useCallback(
    (onStoreChange: () => void) => comp.subscribe(onStoreChange),
    [comp],
  );
  const getSnapshot = useCallback(() => comp.get(), [comp]);
  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
};

/**
 * Read a single atom's value, subscribing the component to it. A `defer(a)`
 * wrapper is read through `useDeferredValue`, so the deferred semantics hold in
 * the hooks world too — the value lags urgent work instead of storming it.
 */
export const useAtomValue = <A,>(atomRead: ReadableAtom<A>): A => {
  const subscribe = useCallback(
    (onStoreChange: () => void) => atomRead.subscribe(onStoreChange),
    [atomRead],
  );
  const getSnapshot = useCallback(() => atomRead.value, [atomRead]);
  // The branch is static per call site: `defer(a)` memoises its wrapper per
  // source, so an atom's deferredness never flips between renders — the hook
  // order holds (same discipline as the interpreter's `atomReader` below).
  if (isDeferred(atomRead)) {
    return useDeferredValue(useSyncExternalStore(subscribe, getSnapshot, getSnapshot));
  }
  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
};

/** A stable setter for a writable atom, supporting both values and updater functions. */
export const useAtomSet = <A,>(atomWrite: Atom<A>): ((value: A | ((prev: A) => A)) => void) =>
  useCallback(
    (value) => {
      if (typeof value === 'function') {
        atomWrite.update(value as (prev: A) => A);
      } else {
        atomWrite.set(value);
      }
    },
    [atomWrite],
  );

/** Read and write a single atom — the `useState` shape, backed by Effect. */
export const useAtom = <A,>(
  atomWrite: Atom<A>,
): readonly [A, (value: A | ((prev: A) => A)) => void] => [
  useAtomValue(atomWrite),
  useAtomSet(atomWrite),
];

export interface ObserveProps<A extends ReactNode> {
  readonly children: (read: Read) => A;
}

/** The render-prop form of {@link observe}, for inline reactive values in JSX. */
export const Observe = <A extends ReactNode>({ children }: ObserveProps<A>): ReactNode =>
  observe(children);

/**
 * The reader the interpreter uses to resolve a yielded atom: read its current
 * value and subscribe this render to it (React's `useSyncExternalStore`). Called
 * once per atom read during the render pass, in the body's stable order. A
 * `defer(a)` wrapper is read through `useDeferredValue` on top — two hooks, but
 * the branch is static per yield site (the wrapper identity is memoised per
 * source and fixed in the body), so the hook order holds across renders.
 */
export const atomReader: Reader = {
  // `atom.subscribe` is a stable reference on the atom, so pass it directly —
  // React must not see a new `subscribe` each render or it re-subscribes every
  // time (and, for an object-valued derived atom, could churn).
  read: (atomRead) =>
    isDeferred(atomRead)
      ? useDeferredValue(
          useSyncExternalStore(atomRead.subscribe, getValue(atomRead), getValue(atomRead)),
        )
      : useSyncExternalStore(atomRead.subscribe, getValue(atomRead), getValue(atomRead)),
};

const snapshots = new WeakMap<ReadableAtom<unknown, unknown, string>, () => unknown>();
/** A per-atom stable `getSnapshot` (`() => atom.value`), memoised by atom identity. */
const getValue = <A,>(atomRead: ReadableAtom<A, unknown, string>): (() => A) => {
  let snapshot = snapshots.get(atomRead) as (() => A) | undefined;
  if (snapshot === undefined) {
    snapshot = () => atomRead.value;
    snapshots.set(atomRead, snapshot);
  }
  return snapshot;
};