citation-work/src/work/EngineContext.tsx

/**
 * Engine + active-document React context.
 *
 * MVP composition root for the UI: one `Engine` instance for the lifetime of
 * the SPA, plus the "what's open in the viewer right now" pointer.
 * `useEngine()` returns the engine; `useActiveDocument()` returns the
 * currently-loaded `{document, representation}` pair, refreshed when the
 * engine emits `DocumentImported` / `DocumentRepresentationGenerated`.
 *
 * Replaces ad-hoc engine wiring inside each component. Per the workplan
 * (T07 note), state lives in a single React context; no Zustand or Redux.
 */

import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useState,
  type ReactNode,
} from "react";
import type { Document, DocumentRepresentation } from "@shared/document";
import type { AnnotationId, DocumentId, SessionId } from "@shared/ids";
import type { Selector } from "@shared/selector";
import {
  attachPersister,
  createEngine,
  engineSnapshotKey,
  restoreFromStorage,
  type Engine,
} from "@engine/index";
import type { SelectionCapture } from "@citation-evidence/evidence-anchor";
import { createPdfByteStore, type PdfByteStore } from "@source/index";
import { useContext as useReactContext } from "react";
import { SessionInternalContext } from "./SessionContextInternal";

/**
 * Legacy single-bucket storage keys, kept for any user landing on a
 * build without sessions. CE-WP-0005 switched persistence to per-session
 * keys (`engineSnapshotKey(sessionId)`); the unscoped keys below are
 * only consulted when no `sessionId` is provided to the provider.
 */
const LEGACY_STORAGE_KEY = "citation-evidence:engine-snapshot:v1";
const ACTIVE_KEY = "citation-evidence:active-document-id:v1";

function storageKeyFor(sessionId: SessionId | null): string {
  return sessionId ? engineSnapshotKey(sessionId) : LEGACY_STORAGE_KEY;
}

function activeDocumentKeyFor(sessionId: SessionId | null): string {
  return sessionId
    ? `citation-evidence:session:${sessionId}:active-document-id:v1`
    : ACTIVE_KEY;
}

/**
 * The pending selection lives in context (not local component state) because
 * the toolbar that consumes it is rendered above the viewer, not inside it.
 * `null` means "no selection waiting for a comment".
 */
export interface PendingSelection {
  readonly capture: SelectionCapture;
  readonly selectors: readonly Selector[];
}

interface EngineContextValue {
  readonly engine: Engine;
  readonly byteStore: PdfByteStore;
  readonly activeDocumentId: DocumentId | null;
  setActiveDocumentId(id: DocumentId | null): void;
  readonly pendingSelection: PendingSelection | null;
  setPendingSelection(pending: PendingSelection | null): void;
  readonly scrollToAnnotationId: AnnotationId | null;
  /** The version counter bumps even when the same id is set twice in a row,
   * so a second click on the same evidence item still triggers a scroll. */
  readonly scrollVersion: number;
  scrollToAnnotation(id: AnnotationId | null): void;
  /**
   * Bumps each time the engine's repos are mutated outside the normal
   * event-emitting service path — currently only on `restoreFromStorage`.
   * Consumers that cache `engine.documents.list()` via `useMemo` add this
   * to their deps so the restored state is reflected on remount.
   */
  readonly engineRevision: number;
}

const EngineContext = createContext<EngineContextValue | null>(null);

interface EngineProviderProps {
  readonly children: ReactNode;
  /** Inject a pre-built engine for tests; production uses the default. */
  readonly engine?: Engine;
  /**
   * Active session id. Drives the per-session storage key for the engine
   * snapshot and the active-document pointer. `null`/omitted falls back
   * to the legacy unscoped keys for back-compat with pre-CE-WP-0005
   * builds.
   *
   * To switch sessions, parents should *re-key* this provider
   * (`<EngineProvider key={sessionId} sessionId={sessionId}>`) so React
   * unmounts the subtree and a fresh engine is created.
   */
  readonly sessionId?: SessionId | null;
}

export function EngineProvider({
  children,
  engine: injected,
  sessionId = null,
}: EngineProviderProps) {
  const engine = useMemo(() => injected ?? createEngine(), [injected]);
  // Prefer the SessionProvider's per-session byte store registry when
  // available; fall back to a provider-local store for tests that mount
  // EngineProvider on its own.
  const sessionCtx = useReactContext(SessionInternalContext);
  const [fallbackByteStore] = useState<PdfByteStore>(() => createPdfByteStore());
  const byteStore =
    sessionId && sessionCtx
      ? sessionCtx.getOrCreateByteStore(sessionId)
      : fallbackByteStore;
  const [activeDocumentId, setActiveDocumentIdState] = useState<DocumentId | null>(null);
  // `restoreFromStorage` writes directly to the engine's repos without
  // firing engine events (by design — see persistence.ts). That means
  // consuming components (CollectionList etc.) wouldn't normally
  // re-render to reflect the restored state. Bumping `engineRevision`
  // after restore is what consumers add to their `useMemo` deps so
  // the restored state shows up on (re-)mount.
  const [engineRevision, setEngineRevision] = useState(0);
  const [pendingSelection, setPendingSelection] = useState<PendingSelection | null>(null);
  const [scrollState, setScrollState] = useState<{ id: AnnotationId | null; version: number }>({
    id: null,
    version: 0,
  });

  const snapshotKey = storageKeyFor(sessionId);
  const activeDocKey = activeDocumentKeyFor(sessionId);

  // Restore from localStorage on first mount, then attach the persister.
  // The injected-engine path skips persistence (tests own their lifecycle).
  useEffect(() => {
    if (injected) return;
    if (typeof globalThis.localStorage === "undefined") return;
    const result = restoreFromStorage(engine, { key: snapshotKey });
    if (result.restored) {
      const saved = globalThis.localStorage.getItem(activeDocKey);
      if (saved && engine.documents.get(saved as DocumentId)) {
        setActiveDocumentIdState(saved as DocumentId);
      }
      // Force a re-render so consumers see the restored repos.
      setEngineRevision((n) => n + 1);
    }
    return attachPersister(engine, { key: snapshotKey });
  }, [engine, injected, snapshotKey, activeDocKey]);

  // Persist the active-document pointer alongside the engine snapshot so a
  // reload lands the user back where they were.
  useEffect(() => {
    if (injected) return;
    if (typeof globalThis.localStorage === "undefined") return;
    if (activeDocumentId) {
      globalThis.localStorage.setItem(activeDocKey, activeDocumentId);
    } else {
      globalThis.localStorage.removeItem(activeDocKey);
    }
  }, [activeDocumentId, injected, activeDocKey]);

  // Switching the active document discards any pending selection — it
  // belongs to the previous document's viewer state.
  const setActiveDocumentId = useCallback((id: DocumentId | null) => {
    setActiveDocumentIdState(id);
    setPendingSelection(null);
    setScrollState((prev) => ({ id: null, version: prev.version + 1 }));
  }, []);

  const scrollToAnnotation = useCallback((id: AnnotationId | null) => {
    setScrollState((prev) => ({ id, version: prev.version + 1 }));
  }, []);

  const value = useMemo<EngineContextValue>(
    () => ({
      engine,
      byteStore,
      activeDocumentId,
      setActiveDocumentId,
      pendingSelection,
      setPendingSelection,
      scrollToAnnotationId: scrollState.id,
      scrollVersion: scrollState.version,
      scrollToAnnotation,
      engineRevision,
    }),
    [
      engine,
      byteStore,
      activeDocumentId,
      setActiveDocumentId,
      pendingSelection,
      scrollState,
      scrollToAnnotation,
      engineRevision,
    ],
  );

  return <EngineContext.Provider value={value}>{children}</EngineContext.Provider>;
}

export function useEngine(): Engine {
  const ctx = useContext(EngineContext);
  if (!ctx) throw new Error("useEngine: missing EngineProvider");
  return ctx.engine;
}

export function usePdfByteStore(): PdfByteStore {
  const ctx = useContext(EngineContext);
  if (!ctx) throw new Error("usePdfByteStore: missing EngineProvider");
  return ctx.byteStore;
}

export function useEngineRevision(): number {
  const ctx = useContext(EngineContext);
  if (!ctx) throw new Error("useEngineRevision: missing EngineProvider");
  return ctx.engineRevision;
}

export function useActiveDocumentId(): {
  readonly id: DocumentId | null;
  setId(id: DocumentId | null): void;
} {
  const ctx = useContext(EngineContext);
  if (!ctx) throw new Error("useActiveDocumentId: missing EngineProvider");
  return { id: ctx.activeDocumentId, setId: ctx.setActiveDocumentId };
}

export function useActiveDocument(): {
  readonly document: Document | null;
  readonly representation: DocumentRepresentation | null;
} {
  const engine = useEngine();
  const { id } = useActiveDocumentId();
  const [tick, setTick] = useState(0);

  // Re-render when documents come and go so list views stay fresh.
  useEffect(() => {
    const off1 = engine.bus.on("DocumentImported", () => setTick((t) => t + 1));
    const off2 = engine.bus.on("DocumentRepresentationGenerated", () => setTick((t) => t + 1));
    return () => {
      off1();
      off2();
    };
  }, [engine]);

  const document = id ? engine.documents.get(id) : null;
  const representation = id
    ? engine.documents.listRepresentations(id).at(-1) ?? null
    : null;

  // `tick` is intentionally read to silence unused-var warnings; the dep
  // chain is via useState so React handles the re-render. We don't actually
  // need to consume the value.
  void tick;

  return { document, representation };
}

/**
 * Subscribe to a single engine event type and trigger a re-render each time
 * it fires. Returns the current monotonic counter — pure state-marker.
 */
export function useEngineEventTick<T extends Parameters<Engine["bus"]["on"]>[0]>(
  type: T,
): number {
  const engine = useEngine();
  const [tick, setTick] = useState(0);
  const bump = useCallback(() => setTick((t) => t + 1), []);
  useEffect(() => engine.bus.on(type, bump), [engine, type, bump]);
  return tick;
}

export function usePendingSelection(): {
  readonly pending: PendingSelection | null;
  set(pending: PendingSelection | null): void;
} {
  const ctx = useContext(EngineContext);
  if (!ctx) throw new Error("usePendingSelection: missing EngineProvider");
  return { pending: ctx.pendingSelection, set: ctx.setPendingSelection };
}

export function useScrollToAnnotation(): {
  readonly id: AnnotationId | null;
  readonly version: number;
  scrollTo(id: AnnotationId | null): void;
} {
  const ctx = useContext(EngineContext);
  if (!ctx) throw new Error("useScrollToAnnotation: missing EngineProvider");
  return {
    id: ctx.scrollToAnnotationId,
    version: ctx.scrollVersion,
    scrollTo: ctx.scrollToAnnotation,
  };
}

/**
 * Track the most-recent `EvidenceItemActivated` event id from the engine
 * bus. Returns `null` until something is activated. UI components that
 * highlight "the active evidence" subscribe via this hook so they don't
 * need to import the binder's active-state machine directly.
 */
export function useLastActivatedEvidence(): import("@shared/ids").EvidenceItemId | null {
  const engine = useEngine();
  const [id, setId] = useState<import("@shared/ids").EvidenceItemId | null>(null);
  useEffect(() => {
    return engine.bus.on("EvidenceItemActivated", (e) => setId(e.evidenceItemId));
  }, [engine]);
  return id;
}