Widget Store

/**
 * Pure React widget model store for Jupyter widgets.
 *
 * This replaces the Backbone.js-based model system from @jupyter-widgets/html-manager
 * with a pure JavaScript store that integrates with React via useSyncExternalStore.
 */

// === Types ===

export interface WidgetModel {
  /** comm_id from the Jupyter protocol */
  id: string;
  /** Widget state (value, min, max, description, etc.) */
  state: Record<string, unknown>;
  /** Binary buffers for data-heavy widgets */
  buffers: ArrayBuffer[];
  /** Model class name, e.g., "IntSliderModel", "AnyModel" */
  modelName: string;
  /** Model module, e.g., "@jupyter-widgets/controls", "anywidget" */
  modelModule: string;
}

type Listener = () => void;
type KeyListener = (value: unknown) => void;
// Anywidgets expect DataView[] so they can access .buffer for the underlying ArrayBuffer
// This matches JupyterLab services which deserializes buffers as DataView[]
type CustomMessageCallback = (
  content: Record<string, unknown>,
  buffers?: DataView[],
) => void;

export interface WidgetStore {
  /** Subscribe to all model changes (for useSyncExternalStore) */
  subscribe: (listener: Listener) => () => void;
  /** Get current models snapshot (for useSyncExternalStore) */
  getSnapshot: () => Map<string, WidgetModel>;
  /** Get a single model by ID */
  getModel: (modelId: string) => WidgetModel | undefined;
  /** Create a new model (on comm_open) */
  createModel: (
    commId: string,
    state: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Update a model's state (on comm_msg with method: "update") */
  updateModel: (
    commId: string,
    statePatch: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Delete a model (on comm_close) */
  deleteModel: (commId: string) => void;
  /** Check if a model was explicitly closed (vs never existed) */
  wasModelClosed: (modelId: string) => boolean;
  /** Subscribe to a specific key on a specific model */
  subscribeToKey: (
    modelId: string,
    key: string,
    callback: KeyListener,
  ) => () => void;
  /** Emit a custom message to listeners for a model */
  emitCustomMessage: (
    commId: string,
    content: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Subscribe to custom messages for a model */
  subscribeToCustomMessage: (
    commId: string,
    callback: CustomMessageCallback,
  ) => () => void;
}

// === IPY_MODEL_ Reference Resolution ===

const IPY_MODEL_PREFIX = "IPY_MODEL_";

/**
 * Check if a value is an IPY_MODEL_ reference string.
 */
export function isModelRef(value: unknown): value is string {
  return typeof value === "string" && value.startsWith(IPY_MODEL_PREFIX);
}

/**
 * Extract the model ID from an IPY_MODEL_ reference.
 */
export function parseModelRef(ref: string): string | null {
  if (ref.startsWith(IPY_MODEL_PREFIX)) {
    return ref.slice(IPY_MODEL_PREFIX.length);
  }
  return null;
}

/**
 * Resolve an IPY_MODEL_ reference to the actual model.
 * If the value is not a reference, returns it unchanged.
 */
export function resolveModelRef(
  value: unknown,
  getModel: (id: string) => WidgetModel | undefined,
): unknown {
  if (isModelRef(value)) {
    const refId = parseModelRef(value);
    return refId ? getModel(refId) : value;
  }
  return value;
}

// === Store Factory ===

/**
 * Create a new widget store instance.
 *
 * The store manages widget models and provides fine-grained subscriptions
 * for reactive updates. It's designed to work with React's useSyncExternalStore.
 */
export function createWidgetStore(): WidgetStore {
  // Internal state - using Map for O(1) lookups
  let models = new Map<string, WidgetModel>();

  // Track explicitly closed models (for distinguishing from "never existed")
  const closedModels = new Set<string>();

  // Global listeners (for useSyncExternalStore)
  const listeners = new Set<Listener>();

  // Key-specific listeners for fine-grained subscriptions
  // Structure: modelId -> key -> Set<callback>
  const keyListeners = new Map<string, Map<string, Set<KeyListener>>>();

  // Custom message listeners
  // Structure: modelId -> Set<callback>
  const customListeners = new Map<string, Set<CustomMessageCallback>>();

  // Buffered custom messages for comm_ids with no listeners yet
  // Structure: commId -> Array<{ content, buffers }>
  const customMessageBuffer = new Map<
    string,
    Array<{ content: Record<string, unknown>; buffers?: DataView[] }>
  >();
  const MAX_BUFFERED_MESSAGES = 1000;

  // Notify all global listeners that something changed
  function emitChange() {
    listeners.forEach((listener) => listener());
  }

  // Notify key-specific listeners for changed keys
  function emitKeyChanges(modelId: string, changedKeys: string[]) {
    const modelListeners = keyListeners.get(modelId);
    if (!modelListeners) return;

    const model = models.get(modelId);
    for (const key of changedKeys) {
      const keyCallbacks = modelListeners.get(key);
      if (keyCallbacks) {
        const value = model?.state[key];
        keyCallbacks.forEach((cb) => cb(value));
      }
    }
  }

  return {
    subscribe(listener: Listener): () => void {
      listeners.add(listener);
      return () => listeners.delete(listener);
    },

    getSnapshot(): Map<string, WidgetModel> {
      return models;
    },

    getModel(modelId: string): WidgetModel | undefined {
      return models.get(modelId);
    },

    createModel(
      commId: string,
      state: Record<string, unknown>,
      buffers: ArrayBuffer[] = [],
    ): void {
      // Handle re-open: remove from closed set if re-opening
      closedModels.delete(commId);

      // Extract model metadata from state
      const modelName = (state._model_name as string) || "UnknownModel";
      const modelModule = (state._model_module as string) || "";

      const model: WidgetModel = {
        id: commId,
        state,
        buffers,
        modelName,
        modelModule,
      };

      // Create new Map to trigger useSyncExternalStore update
      models = new Map(models).set(commId, model);
      emitChange();
    },

    updateModel(
      commId: string,
      statePatch: Record<string, unknown>,
      buffers?: ArrayBuffer[],
    ): void {
      const existing = models.get(commId);
      if (!existing) {
        // Model doesn't exist yet - this can happen if messages arrive out of order
        // In practice, comm_open should always come first
        return;
      }

      // Merge state patch into existing state
      const newState = { ...existing.state, ...statePatch };
      const newModel: WidgetModel = {
        ...existing,
        state: newState,
        buffers: buffers ?? existing.buffers,
      };

      // Create new Map to trigger useSyncExternalStore update
      models = new Map(models).set(commId, newModel);

      // Emit changes
      emitChange();
      emitKeyChanges(commId, Object.keys(statePatch));
    },

    deleteModel(commId: string): void {
      if (!models.has(commId)) return;

      // Track that this model was explicitly closed
      closedModels.add(commId);

      // Create new Map without the deleted model
      models = new Map(models);
      models.delete(commId);

      // Clean up listeners and buffered messages for this model
      keyListeners.delete(commId);
      customListeners.delete(commId);
      customMessageBuffer.delete(commId);

      emitChange();
    },

    subscribeToKey(
      modelId: string,
      key: string,
      callback: KeyListener,
    ): () => void {
      // Ensure model entry exists
      if (!keyListeners.has(modelId)) {
        keyListeners.set(modelId, new Map());
      }
      const modelMap = keyListeners.get(modelId)!;

      // Ensure key entry exists
      if (!modelMap.has(key)) {
        modelMap.set(key, new Set());
      }
      modelMap.get(key)?.add(callback);

      // Return unsubscribe function
      return () => {
        modelMap.get(key)?.delete(callback);

        // Clean up empty sets
        if (modelMap.get(key)?.size === 0) {
          modelMap.delete(key);
        }
        if (modelMap.size === 0) {
          keyListeners.delete(modelId);
        }
      };
    },

    emitCustomMessage(
      commId: string,
      content: Record<string, unknown>,
      buffers?: ArrayBuffer[],
    ): void {
      // Convert ArrayBuffer[] to DataView[] for anywidget compatibility
      // Anywidgets access the underlying buffer via .buffer property
      const dataViewBuffers = buffers?.map((b) =>
        b instanceof DataView ? b : new DataView(b),
      );

      // Always buffer so future subscribers get the history (e.g. a second
      // CanvasWidget subscribing to the same CanvasManagerModel after the
      // first canvas is already mounted and receiving messages).
      if (!customMessageBuffer.has(commId)) {
        customMessageBuffer.set(commId, []);
      }
      const buffer = customMessageBuffer.get(commId)!;
      buffer.push({ content, buffers: dataViewBuffers });
      // Evict oldest messages if over limit
      if (buffer.length > MAX_BUFFERED_MESSAGES) {
        buffer.splice(0, buffer.length - MAX_BUFFERED_MESSAGES);
      }

      // Also deliver to existing subscribers
      const callbacks = customListeners.get(commId);
      if (callbacks && callbacks.size > 0) {
        callbacks.forEach((cb) => cb(content, dataViewBuffers));
      }
    },

    subscribeToCustomMessage(
      commId: string,
      callback: CustomMessageCallback,
    ): () => void {
      // Ensure entry exists
      if (!customListeners.has(commId)) {
        customListeners.set(commId, new Set());
      }
      customListeners.get(commId)?.add(callback);

      // Flush any buffered messages to this new subscriber.
      // Keep the buffer (don't delete) so other subscribers to the same
      // comm_id also receive these messages — e.g. multiple CanvasWidgets
      // subscribing to one CanvasManagerModel.
      const buffered = customMessageBuffer.get(commId);
      if (buffered && buffered.length > 0) {
        for (const msg of buffered) {
          callback(msg.content, msg.buffers);
        }
      }

      // Return unsubscribe function
      return () => {
        customListeners.get(commId)?.delete(callback);

        // Clean up empty sets
        if (customListeners.get(commId)?.size === 0) {
          customListeners.delete(commId);
        }
      };
    },

    wasModelClosed(modelId: string): boolean {
      return closedModels.has(modelId);
    },
  };
}

"use client";

/**
 * React context and hooks for the widget model store.
 *
 * Provides:
 * - WidgetStoreProvider: Wrap your app to enable widget support
 * - useWidgetStore: Access the store context (nullable)
 * - useWidgetModels: Subscribe to all models
 * - useWidgetModel: Subscribe to a single model
 * - useWidgetModelValue: Subscribe to a single key (finest granularity)
 */

import {
  createContext,
  type ReactNode,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useSyncExternalStore,
} from "react";
import { createCanvasManagerRouter } from "./canvas-manager-subscriptions";
import { createLinkManager } from "./link-subscriptions";
import {
  type JupyterCommMessage,
  type SendMessage,
  useCommRouter,
} from "./use-comm-router";
import {
  createWidgetStore,
  resolveModelRef,
  type WidgetModel,
  type WidgetStore,
} from "./widget-store";

// === Context Types ===

interface WidgetStoreContextValue {
  store: WidgetStore;
  /** Handle incoming Jupyter comm messages */
  handleMessage: (msg: JupyterCommMessage) => void;
  /** Raw send function (prefer sendUpdate/sendCustom for specific cases) */
  sendMessage: SendMessage;
  /** Send a state update to the kernel */
  sendUpdate: (
    commId: string,
    state: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Send a custom message to the kernel */
  sendCustom: (
    commId: string,
    content: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Close a comm channel */
  closeComm: (commId: string) => void;
}

// === Context ===

// Export context for use in isolated iframe widget provider
export const WidgetStoreContext = createContext<WidgetStoreContextValue | null>(
  null,
);

// === Provider ===

interface WidgetStoreProviderProps {
  children: ReactNode;
  /** Function to send messages back to the kernel (for widget interactions) */
  sendMessage?: SendMessage;
}

/**
 * Provider component for the widget store.
 *
 * Wrap your app (or the part that needs widgets) with this provider.
 * Pass a sendMessage function to enable widget interactions back to the kernel.
 */
export function WidgetStoreProvider({
  children,
  sendMessage = () => {},
}: WidgetStoreProviderProps) {
  // Create store once and keep it stable across renders
  const storeRef = useRef<WidgetStore | null>(null);
  if (!storeRef.current) {
    storeRef.current = createWidgetStore();
  }
  const store = storeRef.current;

  // Use the comm router hook for message handling
  const { handleMessage, sendUpdate, sendCustom, closeComm } = useCommRouter({
    sendMessage,
    store,
  });

  // Manage link subscriptions (jslink/jsdlink) at the store level.
  // Headless widgets like LinkModel have _view_name: null and won't be
  // in any container's children, so they need store-level subscriptions.
  // Pass sendUpdate so linked values sync back to the Python kernel.
  useEffect(() => createLinkManager(store, sendUpdate), [store, sendUpdate]);
  useEffect(() => createCanvasManagerRouter(store), [store]);

  const value = useMemo(
    () => ({
      store,
      handleMessage,
      sendMessage,
      sendUpdate,
      sendCustom,
      closeComm,
    }),
    [store, handleMessage, sendMessage, sendUpdate, sendCustom, closeComm],
  );

  return (
    <WidgetStoreContext.Provider value={value}>
      {children}
    </WidgetStoreContext.Provider>
  );
}

// === Hooks ===

/**
 * Access the widget store context.
 * Returns null if used outside of WidgetStoreProvider.
 */
export function useWidgetStore(): WidgetStoreContextValue | null {
  return useContext(WidgetStoreContext);
}

/**
 * Access the widget store context, throwing if not available.
 * Use this when you know you're inside a WidgetStoreProvider.
 */
export function useWidgetStoreRequired(): WidgetStoreContextValue {
  const ctx = useContext(WidgetStoreContext);
  if (!ctx) {
    throw new Error(
      "useWidgetStoreRequired must be used within WidgetStoreProvider",
    );
  }
  return ctx;
}

/**
 * Subscribe to all widget models.
 * Re-renders when any model is added, updated, or removed.
 */
export function useWidgetModels(): Map<string, WidgetModel> {
  const { store } = useWidgetStoreRequired();

  return useSyncExternalStore(
    store.subscribe,
    store.getSnapshot,
    store.getSnapshot, // SSR snapshot (same as client)
  );
}

/**
 * Subscribe to a specific widget model.
 * Re-renders when that model's state changes.
 */
export function useWidgetModel(modelId: string): WidgetModel | undefined {
  const { store } = useWidgetStoreRequired();

  const subscribe = useCallback(
    (callback: () => void) => store.subscribe(callback),
    [store],
  );

  const getSnapshot = useCallback(
    () => store.getModel(modelId),
    [store, modelId],
  );

  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
}

/**
 * Subscribe to a specific key in a widget model's state.
 * This is the finest granularity - only re-renders when that specific key changes.
 *
 * @example
 * const value = useWidgetModelValue<number>(modelId, 'value');
 * const description = useWidgetModelValue<string>(modelId, 'description');
 */
export function useWidgetModelValue<T = unknown>(
  modelId: string,
  key: string,
): T | undefined {
  const { store } = useWidgetStoreRequired();

  const subscribe = useCallback(
    (callback: () => void) => store.subscribeToKey(modelId, key, callback),
    [store, modelId, key],
  );

  const getSnapshot = useCallback(
    () => store.getModel(modelId)?.state[key] as T | undefined,
    [store, modelId, key],
  );

  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
}

/**
 * Get a value from a widget model, resolving IPY_MODEL_ references.
 * If the value is an IPY_MODEL_<id> reference, returns the referenced model.
 *
 * @example
 * // If state.layout is "IPY_MODEL_abc123", this returns the LayoutModel
 * const layout = useResolvedModelValue(modelId, 'layout');
 */
export function useResolvedModelValue<T = unknown>(
  modelId: string,
  key: string,
): T | WidgetModel | undefined {
  const { store } = useWidgetStoreRequired();
  const value = useWidgetModelValue(modelId, key);

  // Resolve IPY_MODEL_ reference if applicable
  const resolved = resolveModelRef(value, (id) => store.getModel(id));

  return resolved as T | WidgetModel | undefined;
}

/**
 * Check if a widget model was explicitly closed (e.g., tqdm with leave=False).
 * Returns true if the model was closed, false if it never existed or is active.
 *
 * Note: This hook is not reactive on its own. It relies on being used alongside
 * useWidgetModel, which subscribes to store changes and triggers re-renders.
 */
export function useWasWidgetClosed(modelId: string): boolean {
  const { store } = useWidgetStoreRequired();
  return store.wasModelClosed(modelId);
}

// Re-export store-level managers for non-React integrations (e.g. iframe isolation)
export { createCanvasManagerRouter } from "./canvas-manager-subscriptions";
export { createLinkManager } from "./link-subscriptions";
export type {
  JupyterCommMessage,
  JupyterMessageHeader,
  SendMessage,
} from "./use-comm-router";
// Re-export types from use-comm-router
export { useCommRouter } from "./use-comm-router";
export type { WidgetModel, WidgetStore } from "./widget-store";
// Re-export types and utilities from widget-store
export { isModelRef, parseModelRef, resolveModelRef } from "./widget-store";

/**
 * Comm Protocol Router Hook
 *
 * Handles routing of Jupyter comm messages between the kernel and widget store.
 * Provides both inbound message handling and outbound message construction.
 *
 * Message format follows the Jupyter messaging protocol specification and is
 * compatible with strongly-typed backends (Rust, Go) that require explicit fields.
 *
 * @see https://jupyter-widgets.readthedocs.io/en/latest/examples/Widget%20Low%20Level.html
 * @see https://jupyter-client.readthedocs.io/en/latest/messaging.html
 */

import { useCallback, useEffect, useRef } from "react";
import { applyBufferPaths } from "./buffer-utils";
import type { WidgetStore } from "./widget-store";

// === Message Types ===

/**
 * Jupyter message header.
 * Required fields vary between incoming (may have fewer) and outgoing (should have all).
 */
export interface JupyterMessageHeader {
  msg_id: string;
  msg_type: string;
  username?: string;
  session?: string;
  date?: string;
  version?: string;
}

/**
 * Full header with all fields required (for outgoing messages).
 */
interface FullJupyterMessageHeader {
  msg_id: string;
  msg_type: string;
  username: string;
  session: string;
  date: string;
  version: string;
}

/**
 * Jupyter comm message.
 * Some fields are optional for incoming messages but should be set for outgoing.
 */
export interface JupyterCommMessage {
  header: JupyterMessageHeader;
  /** Should be null for outgoing messages (not undefined or empty object) */
  parent_header?: JupyterMessageHeader | null;
  metadata?: Record<string, unknown>;
  content: {
    comm_id?: string;
    target_name?: string;
    data?: {
      state?: Record<string, unknown>;
      method?: string;
      content?: Record<string, unknown>;
      buffer_paths?: string[][];
      [key: string]: unknown;
    };
  };
  buffers?: ArrayBuffer[];
  channel?: string | null;
}

/**
 * Outgoing message with all fields populated for protocol compliance.
 * Use this type for messages being sent to strongly-typed backends.
 */
interface OutgoingJupyterCommMessage {
  header: FullJupyterMessageHeader;
  parent_header: null;
  metadata: Record<string, unknown>;
  content: {
    comm_id: string;
    data?: {
      state?: Record<string, unknown>;
      method?: string;
      content?: Record<string, unknown>;
      buffer_paths: string[][];
    };
  };
  buffers: ArrayBuffer[];
  channel: string;
}

/**
 * Function type for sending messages to the kernel.
 */
export type SendMessage = (msg: JupyterCommMessage) => void;

// === Hook Types ===

export interface UseCommRouterOptions {
  /** Function to send messages to the kernel */
  sendMessage: SendMessage;
  /** Widget store instance */
  store: WidgetStore;
  /** Optional username for message headers (default: "frontend") */
  username?: string;
}

export interface UseCommRouterReturn {
  /** Handle incoming Jupyter comm messages */
  handleMessage: (msg: JupyterCommMessage) => void;
  /** Send a state update to the kernel */
  sendUpdate: (
    commId: string,
    state: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Send a custom message to the kernel */
  sendCustom: (
    commId: string,
    content: Record<string, unknown>,
    buffers?: ArrayBuffer[],
  ) => void;
  /** Close a comm channel */
  closeComm: (commId: string) => void;
}

// === Message Construction Helpers ===

// Session ID for this frontend instance (stable across messages)
const SESSION_ID = crypto.randomUUID();

/**
 * Create a complete Jupyter message header with all fields.
 * All fields are required for compatibility with strongly-typed backends.
 */
function createHeader(
  msgType: string,
  username: string,
): FullJupyterMessageHeader {
  return {
    msg_id: crypto.randomUUID(),
    msg_type: msgType,
    username,
    session: SESSION_ID,
    date: new Date().toISOString(),
    version: "5.3",
  };
}

/**
 * Create a comm_msg for state updates.
 * Includes all required fields for Jupyter protocol compliance.
 */
function createUpdateMessage(
  commId: string,
  state: Record<string, unknown>,
  buffers: ArrayBuffer[] | undefined,
  username: string,
): OutgoingJupyterCommMessage {
  return {
    header: createHeader("comm_msg", username),
    parent_header: null,
    metadata: {},
    content: {
      comm_id: commId,
      data: {
        method: "update",
        state,
        buffer_paths: [],
      },
    },
    buffers: buffers ?? [],
    channel: "shell",
  };
}

/**
 * Create a comm_msg for custom messages.
 * Includes all required fields for Jupyter protocol compliance.
 */
function createCustomMessage(
  commId: string,
  content: Record<string, unknown>,
  buffers: ArrayBuffer[] | undefined,
  username: string,
): OutgoingJupyterCommMessage {
  return {
    header: createHeader("comm_msg", username),
    parent_header: null,
    metadata: {},
    content: {
      comm_id: commId,
      data: {
        method: "custom",
        content,
        buffer_paths: [],
      },
    },
    buffers: buffers ?? [],
    channel: "shell",
  };
}

/**
 * Create a comm_close message.
 * Includes all required fields for Jupyter protocol compliance.
 */
function createCloseMessage(
  commId: string,
  username: string,
): OutgoingJupyterCommMessage {
  return {
    header: createHeader("comm_close", username),
    parent_header: null,
    metadata: {},
    content: {
      comm_id: commId,
    },
    buffers: [],
    channel: "shell",
  };
}

// === Hook Implementation ===

/**
 * Hook for routing Jupyter comm protocol messages.
 *
 * Handles:
 * - Inbound: comm_open, comm_msg (update/custom), comm_close
 * - Outbound: sendUpdate, sendCustom, closeComm
 *
 * @example
 * const { handleMessage, sendUpdate, sendCustom, closeComm } = useCommRouter({
 *   sendMessage: (msg) => kernel.send(msg),
 *   store: widgetStore,
 * });
 *
 * // Route incoming messages
 * kernel.onMessage((msg) => handleMessage(msg));
 *
 * // Send updates back to kernel
 * sendUpdate(commId, { value: 42 });
 */
export function useCommRouter({
  sendMessage,
  store,
  username = "frontend",
}: UseCommRouterOptions): UseCommRouterReturn {
  // Refs for values that need to be fresh but shouldn't cause function identity changes.
  // This ensures sendUpdate, sendCustom, etc. are stable across renders, preventing
  // unnecessary effect re-runs in consumer widgets (see PR #137 pattern).
  const sendMessageRef = useRef(sendMessage);
  const storeRef = useRef(store);
  const usernameRef = useRef(username);

  // Keep refs up-to-date without changing function identities
  useEffect(() => {
    sendMessageRef.current = sendMessage;
    storeRef.current = store;
    usernameRef.current = username;
  });

  /**
   * Handle incoming Jupyter comm messages.
   * Routes to appropriate store methods based on message type.
   */
  const handleMessage = useCallback((msg: JupyterCommMessage) => {
    const msgType = msg.header.msg_type;
    const commId = msg.content.comm_id;

    if (!commId) return;

    switch (msgType) {
      case "comm_open": {
        // Get state and apply buffer paths if present
        let state = msg.content.data?.state || {};
        const bufferPaths = msg.content.data?.buffer_paths;

        if (bufferPaths && msg.buffers?.length) {
          state = { ...state };
          applyBufferPaths(state, bufferPaths, msg.buffers);
        }

        storeRef.current.createModel(commId, state, msg.buffers);
        break;
      }

      case "comm_msg": {
        const data = msg.content.data;
        const method = data?.method;

        if (method === "update" && data?.state) {
          // Apply buffer paths to state update
          let state = data.state;
          const bufferPaths = data.buffer_paths;

          if (bufferPaths && msg.buffers?.length) {
            state = { ...state };
            applyBufferPaths(state, bufferPaths, msg.buffers);
          }

          storeRef.current.updateModel(commId, state, msg.buffers);
        } else if (method === "custom") {
          // Dispatch custom message to widget handlers
          const content = (data?.content as Record<string, unknown>) || {};
          storeRef.current.emitCustomMessage(commId, content, msg.buffers);
        }
        break;
      }

      case "comm_close": {
        storeRef.current.deleteModel(commId);
        break;
      }
    }
  }, []);

  /**
   * Send a state update to the kernel.
   * Also applies optimistic update to local store for immediate UI response.
   */
  const sendUpdate = useCallback(
    (
      commId: string,
      state: Record<string, unknown>,
      buffers?: ArrayBuffer[],
    ) => {
      // Optimistic update: apply locally first for responsive UI
      storeRef.current.updateModel(commId, state, buffers);
      // Then send to kernel
      sendMessageRef.current(
        createUpdateMessage(commId, state, buffers, usernameRef.current),
      );
    },
    [],
  );

  /**
   * Send a custom message to the kernel.
   */
  const sendCustom = useCallback(
    (
      commId: string,
      content: Record<string, unknown>,
      buffers?: ArrayBuffer[],
    ) => {
      sendMessageRef.current(
        createCustomMessage(commId, content, buffers, usernameRef.current),
      );
    },
    [],
  );

  /**
   * Close a comm channel.
   * Sends comm_close to kernel and removes model from store.
   */
  const closeComm = useCallback((commId: string) => {
    sendMessageRef.current(createCloseMessage(commId, usernameRef.current));
    storeRef.current.deleteModel(commId);
  }, []);

  return {
    handleMessage,
    sendUpdate,
    sendCustom,
    closeComm,
  };
}

/**
 * Buffer path utilities for Jupyter widget protocol.
 *
 * Handles the buffer_paths field in comm messages which specifies
 * where binary buffers should be inserted into the JSON message data.
 *
 * @see https://jupyter-widgets.readthedocs.io/en/latest/examples/Widget%20Low%20Level.html
 */

/**
 * Apply buffers to message data at specified paths.
 *
 * The Jupyter widget protocol sends binary data separately from JSON.
 * The buffer_paths field specifies where each buffer should be inserted.
 *
 * @example
 * // buffer_paths: [["state", "data"], ["content", "blob"]]
 * // buffers: [ArrayBuffer1, ArrayBuffer2]
 * // data: { state: {}, content: {} }
 * // Result: { state: { data: ArrayBuffer1 }, content: { blob: ArrayBuffer2 } }
 *
 * @param data - The message data object to modify (mutated in place)
 * @param bufferPaths - Array of paths, each path is an array of keys
 * @param buffers - Array of buffers to insert at the corresponding paths
 * @returns The modified data object (same reference as input)
 */
export function applyBufferPaths(
  data: Record<string, unknown>,
  bufferPaths: string[][] | undefined,
  buffers: ArrayBuffer[] | undefined,
): Record<string, unknown> {
  if (!bufferPaths || !buffers || bufferPaths.length === 0) {
    return data;
  }

  for (let i = 0; i < bufferPaths.length && i < buffers.length; i++) {
    const path = bufferPaths[i];
    const buffer = buffers[i];

    if (path.length === 0) continue;

    // Navigate to the parent object
    let current: Record<string, unknown> = data;
    for (let j = 0; j < path.length - 1; j++) {
      const key = path[j];
      if (current[key] === undefined || current[key] === null) {
        current[key] = {};
      }
      current = current[key] as Record<string, unknown>;
    }

    // Set the buffer at the final key
    const finalKey = path[path.length - 1];
    current[finalKey] = buffer;
  }

  return data;
}

/**
 * Extract buffers from message data at specified paths.
 *
 * This is the reverse of applyBufferPaths - used when sending
 * messages from the frontend to the kernel.
 *
 * @param data - The message data object (mutated in place, buffer values set to null)
 * @param bufferPaths - Array of paths where buffers should be extracted from
 * @returns Array of extracted buffers (or undefined for paths that don't exist)
 */
export function extractBuffers(
  data: Record<string, unknown>,
  bufferPaths: string[][] | undefined,
): ArrayBuffer[] {
  if (!bufferPaths || bufferPaths.length === 0) {
    return [];
  }

  const buffers: ArrayBuffer[] = [];

  for (const path of bufferPaths) {
    if (path.length === 0) {
      buffers.push(new ArrayBuffer(0));
      continue;
    }

    // Navigate to the parent object
    let current: Record<string, unknown> = data;
    let found = true;

    for (let j = 0; j < path.length - 1; j++) {
      const key = path[j];
      if (
        current[key] === undefined ||
        current[key] === null ||
        typeof current[key] !== "object"
      ) {
        found = false;
        break;
      }
      current = current[key] as Record<string, unknown>;
    }

    if (found) {
      const finalKey = path[path.length - 1];
      const value = current[finalKey];

      if (value instanceof ArrayBuffer) {
        buffers.push(value);
        // Replace with null in the data (per protocol spec)
        current[finalKey] = null;
      } else if (
        ArrayBuffer.isView(value) &&
        value.buffer instanceof ArrayBuffer
      ) {
        // Handle typed arrays (Uint8Array, etc.) with ArrayBuffer backing
        buffers.push(
          value.buffer.slice(
            value.byteOffset,
            value.byteOffset + value.byteLength,
          ),
        );
        current[finalKey] = null;
      } else {
        // Path exists but value is not a buffer
        buffers.push(new ArrayBuffer(0));
      }
    } else {
      // Path not found
      buffers.push(new ArrayBuffer(0));
    }
  }

  return buffers;
}

/**
 * Convert an ArrayBuffer or Uint8Array to a base64 string.
 */
export function arrayBufferToBase64(buffer: ArrayBuffer | Uint8Array): string {
  const bytes = buffer instanceof Uint8Array ? buffer : new Uint8Array(buffer);
  let binary = "";
  for (let i = 0; i < bytes.byteLength; i++) {
    binary += String.fromCharCode(bytes[i]);
  }
  return btoa(binary);
}

/**
 * Build a media src URL from a widget value that may be a string or binary data.
 *
 * Handles all value types sent by the Jupyter widget protocol:
 * - Binary data (ArrayBuffer/Uint8Array) from from_url() or from_file()
 * - Data URLs, HTTP URLs, or absolute paths (passed through)
 * - Plain base64 strings (wrapped in a data URL)
 * - Falsy values (returns undefined)
 */
export function buildMediaSrc(
  value: string | ArrayBuffer | Uint8Array | null | undefined,
  mediaType: string,
  format: string,
): string | undefined {
  if (!value) return undefined;

  if (value instanceof ArrayBuffer || value instanceof Uint8Array) {
    const base64 = arrayBufferToBase64(value);
    return `data:${mediaType}/${format};base64,${base64}`;
  }

  if (typeof value === "string") {
    if (
      value.startsWith("data:") ||
      value.startsWith("http://") ||
      value.startsWith("https://") ||
      value.startsWith("/")
    ) {
      return value;
    }
    return `data:${mediaType}/${format};base64,${value}`;
  }

  return undefined;
}

/**
 * Find all ArrayBuffer values in an object and return their paths.
 *
 * Useful for automatically detecting buffer paths when sending data.
 *
 * @param data - The data object to scan
 * @param prefix - Current path prefix (for recursion)
 * @returns Array of paths to ArrayBuffer values
 */
export function findBufferPaths(
  data: Record<string, unknown>,
  prefix: string[] = [],
): string[][] {
  const paths: string[][] = [];

  for (const [key, value] of Object.entries(data)) {
    const currentPath = [...prefix, key];

    if (value instanceof ArrayBuffer || ArrayBuffer.isView(value)) {
      paths.push(currentPath);
    } else if (
      value !== null &&
      typeof value === "object" &&
      !Array.isArray(value)
    ) {
      // Recurse into nested objects
      paths.push(
        ...findBufferPaths(value as Record<string, unknown>, currentPath),
      );
    }
  }

  return paths;
}

/**
 * Store-level link subscriptions for frontend-only property synchronization.
 *
 * Manages LinkModel (jslink) and DirectionalLinkModel (jsdlink) at the
 * store level so links work without React component mounting. This is
 * critical because link widgets are headless (_view_name: null) and won't
 * appear in any container widget's children.
 *
 * Usage:
 *   const cleanup = createLinkManager(store);
 *   // ... later, to tear down all links:
 *   cleanup();
 *
 * WidgetStoreProvider calls this automatically. For non-React integrations
 * (e.g. iframe isolation), call createLinkManager directly.
 */

import { parseModelRef, type WidgetStore } from "./widget-store";

/**
 * Parse a link source/target tuple from widget state.
 * The state arrives as: ["IPY_MODEL_<id>", "attribute_name"]
 * Returns [modelId, attrName] or null if malformed.
 */
function parseLinkTarget(tuple: unknown): [string, string] | null {
  if (
    !Array.isArray(tuple) ||
    tuple.length !== 2 ||
    typeof tuple[0] !== "string" ||
    typeof tuple[1] !== "string"
  ) {
    return null;
  }
  const modelId = parseModelRef(tuple[0]);
  if (!modelId) return null;
  return [modelId, tuple[1]];
}

/** Type for the sendUpdate function that syncs state to the kernel */
type SendUpdate = (commId: string, state: Record<string, unknown>) => void;

/**
 * Set up a one-way property subscription (source → target).
 * Returns a cleanup function to tear down the subscription.
 */
function setupDirectionalLink(
  store: WidgetStore,
  linkModelId: string,
  sendUpdate: SendUpdate,
): () => void {
  let keyUnsub: (() => void) | undefined;
  let globalUnsub: (() => void) | undefined;
  let isSetUp = false;

  function trySetup() {
    if (isSetUp) return;

    const linkModel = store.getModel(linkModelId);
    if (!linkModel) return;

    const src = parseLinkTarget(linkModel.state.source);
    const tgt = parseLinkTarget(linkModel.state.target);
    if (!src || !tgt) return;

    const [sourceModelId, sourceAttr] = src;
    const [targetModelId, targetAttr] = tgt;

    if (!store.getModel(sourceModelId) || !store.getModel(targetModelId)) {
      return;
    }
    isSetUp = true;

    // Initial sync: read source value, write to target
    const sourceModel = store.getModel(sourceModelId);
    if (sourceModel) {
      const currentValue = sourceModel.state[sourceAttr];
      if (currentValue !== undefined) {
        sendUpdate(targetModelId, { [targetAttr]: currentValue });
      }
    }

    // Subscribe to source changes, propagate to target
    keyUnsub = store.subscribeToKey(sourceModelId, sourceAttr, (newValue) => {
      sendUpdate(targetModelId, { [targetAttr]: newValue });
    });

    // Clean up global listener once setup is complete
    if (globalUnsub) {
      globalUnsub();
      globalUnsub = undefined;
    }
  }

  trySetup();

  // If source/target models aren't ready yet, wait for them
  if (!isSetUp) {
    globalUnsub = store.subscribe(() => trySetup());
  }

  return () => {
    globalUnsub?.();
    keyUnsub?.();
  };
}

/**
 * Set up a bidirectional property subscription (source ↔ target).
 * Uses a synchronous guard to prevent infinite loops since store
 * subscriptions fire synchronously during updateModel.
 * Returns a cleanup function to tear down the subscriptions.
 */
function setupBidirectionalLink(
  store: WidgetStore,
  linkModelId: string,
  sendUpdate: SendUpdate,
): () => void {
  const keyUnsubs: (() => void)[] = [];
  let globalUnsub: (() => void) | undefined;
  let isSetUp = false;
  let isSyncing = false;

  function trySetup() {
    if (isSetUp) return;

    const linkModel = store.getModel(linkModelId);
    if (!linkModel) return;

    const src = parseLinkTarget(linkModel.state.source);
    const tgt = parseLinkTarget(linkModel.state.target);
    if (!src || !tgt) return;

    const [sourceModelId, sourceAttr] = src;
    const [targetModelId, targetAttr] = tgt;

    if (!store.getModel(sourceModelId) || !store.getModel(targetModelId)) {
      return;
    }
    isSetUp = true;

    // Initial sync: source → target
    const sourceModel = store.getModel(sourceModelId);
    if (sourceModel) {
      const currentValue = sourceModel.state[sourceAttr];
      if (currentValue !== undefined) {
        isSyncing = true;
        sendUpdate(targetModelId, { [targetAttr]: currentValue });
        isSyncing = false;
      }
    }

    // Source → Target
    keyUnsubs.push(
      store.subscribeToKey(sourceModelId, sourceAttr, (newValue) => {
        if (isSyncing) return;
        isSyncing = true;
        sendUpdate(targetModelId, { [targetAttr]: newValue });
        isSyncing = false;
      }),
    );

    // Target → Source
    keyUnsubs.push(
      store.subscribeToKey(targetModelId, targetAttr, (newValue) => {
        if (isSyncing) return;
        isSyncing = true;
        sendUpdate(sourceModelId, { [sourceAttr]: newValue });
        isSyncing = false;
      }),
    );

    // Clean up global listener once setup is complete
    if (globalUnsub) {
      globalUnsub();
      globalUnsub = undefined;
    }
  }

  trySetup();

  // If source/target models aren't ready yet, wait for them
  if (!isSetUp) {
    globalUnsub = store.subscribe(() => trySetup());
  }

  return () => {
    globalUnsub?.();
    keyUnsubs.forEach((unsub) => unsub());
  };
}

/**
 * Create a link manager that monitors the store for LinkModel and
 * DirectionalLinkModel widgets and manages their property subscriptions.
 *
 * Returns a cleanup function that tears down all active link subscriptions.
 *
 * Called automatically by WidgetStoreProvider. For non-React integrations
 * (e.g. iframe isolation), call this directly after creating the store.
 *
 * @param store - The widget store instance
 * @param sendUpdate - Function to send state updates to the kernel. This ensures
 *   linked widget values are synced back to Python, not just updated in the frontend.
 */
export function createLinkManager(
  store: WidgetStore,
  sendUpdate: SendUpdate,
): () => void {
  const activeLinks = new Map<string, () => void>();
  let lastSize = -1;

  function scan() {
    const models = store.getSnapshot();

    // Only do a full scan when models are added or removed.
    // State updates (e.g. slider drags) don't change the map size.
    if (models.size === lastSize) return;
    lastSize = models.size;

    // Set up new links
    models.forEach((model, id) => {
      if (activeLinks.has(id)) return;

      if (model.modelName === "DirectionalLinkModel") {
        activeLinks.set(id, setupDirectionalLink(store, id, sendUpdate));
      } else if (model.modelName === "LinkModel") {
        activeLinks.set(id, setupBidirectionalLink(store, id, sendUpdate));
      }
    });

    // Clean up removed links
    for (const [id, cleanup] of activeLinks) {
      if (!models.has(id)) {
        cleanup();
        activeLinks.delete(id);
      }
    }
  }

  const unsubscribe = store.subscribe(scan);
  scan();

  return () => {
    unsubscribe();
    activeLinks.forEach((cleanup) => cleanup());
    activeLinks.clear();
  };
}

/**
 * Store-level routing for ipycanvas CanvasManagerModel.
 *
 * CanvasManagerModel is a headless widget (_view_name: null) that receives
 * ALL drawing commands from Python. This module subscribes to each manager's
 * custom messages, parses switchCanvas targets, and re-emits to each target
 * canvas's comm_id — isolating canvases from each other.
 *
 * Same pattern as link-subscriptions.ts for LinkModel/DirectionalLinkModel.
 *
 * Usage:
 *   const cleanup = createCanvasManagerRouter(store);
 *   // ... later, to tear down all routing:
 *   cleanup();
 *
 * WidgetStoreProvider calls this automatically. For non-React integrations
 * (e.g. iframe isolation), call createCanvasManagerRouter directly.
 */

import type { WidgetStore } from "./widget-store";

// ipycanvas drawing command names indexed by protocol number.
// Duplicated from ipycanvas-commands.ts to keep the router self-contained
// within the widget-store package. Must match the Python-side enum in ipycanvas.
const COMMANDS = [
  "fillRect",
  "strokeRect",
  "fillRects",
  "strokeRects",
  "clearRect",
  "fillArc",
  "fillCircle",
  "strokeArc",
  "strokeCircle",
  "fillArcs",
  "strokeArcs",
  "fillCircles",
  "strokeCircles",
  "strokeLine",
  "beginPath",
  "closePath",
  "stroke",
  "strokePath",
  "fillPath",
  "fill",
  "moveTo",
  "lineTo",
  "rect",
  "arc",
  "ellipse",
  "arcTo",
  "quadraticCurveTo",
  "bezierCurveTo",
  "fillText",
  "strokeText",
  "setLineDash",
  "drawImage",
  "putImageData",
  "clip",
  "save",
  "restore",
  "translate",
  "rotate",
  "scale",
  "transform",
  "setTransform",
  "resetTransform",
  "set",
  "clear",
  "sleep",
  "fillPolygon",
  "strokePolygon",
  "strokeLines",
  "fillPolygons",
  "strokePolygons",
  "strokeLineSegments",
  "fillStyledRects",
  "strokeStyledRects",
  "fillStyledCircles",
  "strokeStyledCircles",
  "fillStyledArcs",
  "strokeStyledArcs",
  "fillStyledPolygons",
  "strokeStyledPolygons",
  "strokeStyledLineSegments",
  "switchCanvas",
] as const;

/**
 * Convert a DataView to a TypedArray based on dtype metadata.
 * Inlined here to avoid depending on ipycanvas-commands.ts,
 * keeping the router self-contained within the widget-store package.
 */
function getTypedArray(
  dataview: DataView,
  metadata: { dtype: string },
): ArrayBufferView {
  const buffer = dataview.buffer;
  switch (metadata.dtype) {
    case "int8":
      return new Int8Array(buffer);
    case "uint8":
      return new Uint8Array(buffer);
    case "int16":
      return new Int16Array(buffer);
    case "uint16":
      return new Uint16Array(buffer);
    case "int32":
      return new Int32Array(buffer);
    case "uint32":
      return new Uint32Array(buffer);
    case "float32":
      return new Float32Array(buffer);
    case "float64":
      return new Float64Array(buffer);
    default:
      return new Uint8Array(buffer);
  }
}

/**
 * Walk a command structure and collect switchCanvas target IDs.
 * Calls setTarget as a side effect so subsequent messages without
 * switchCanvas route to the last known target.
 */
function collectSwitchCanvasTargets(
  commands: unknown,
  targets: Set<string>,
  setTarget: (id: string) => void,
): void {
  if (!Array.isArray(commands) || commands.length === 0) return;

  if (Array.isArray(commands[0])) {
    for (const sub of commands) {
      collectSwitchCanvasTargets(sub, targets, setTarget);
    }
  } else {
    const cmdIndex = commands[0] as number;
    if (COMMANDS[cmdIndex] === "switchCanvas") {
      const args = commands[1] as string[] | undefined;
      const ref = args?.[0] ?? "";
      const targetId = ref.startsWith("IPY_MODEL_") ? ref.slice(10) : ref;
      setTarget(targetId);
      targets.add(targetId);
    }
  }
}

/**
 * Set up message routing for a single CanvasManagerModel.
 * Returns a cleanup function to tear down the subscription.
 */
function setupManagerRouting(
  store: WidgetStore,
  managerId: string,
): () => void {
  let currentTarget: string | null = null;

  return store.subscribeToCustomMessage(managerId, (content, buffers) => {
    if (!buffers || buffers.length === 0) return;

    try {
      const metadata = content as { dtype: string };
      const typedArray = getTypedArray(buffers[0], metadata);
      const jsonStr = new TextDecoder("utf-8").decode(typedArray);
      const commands = JSON.parse(jsonStr);

      const targets = new Set<string>();
      collectSwitchCanvasTargets(commands, targets, (id) => {
        currentTarget = id;
      });

      if (targets.size === 0 && currentTarget) {
        targets.add(currentTarget);
      }

      const rawBuffers = buffers.map((dv) => dv.buffer as ArrayBuffer);
      for (const targetId of targets) {
        store.emitCustomMessage(targetId, content, rawBuffers);
      }
    } catch (err) {
      console.warn("[ipycanvas] Manager routing error:", err);
    }
  });
}

/**
 * Create a canvas manager router that monitors the store for
 * CanvasManagerModel widgets and routes their messages to target canvases.
 *
 * Returns a cleanup function that tears down all active routing.
 *
 * Called automatically by WidgetStoreProvider. For non-React integrations
 * (e.g. iframe isolation), call this directly after creating the store.
 */
export function createCanvasManagerRouter(store: WidgetStore): () => void {
  const activeRoutes = new Map<string, () => void>();
  let lastSize = -1;

  function scan() {
    const models = store.getSnapshot();

    if (models.size === lastSize) return;
    lastSize = models.size;

    // Find CanvasModel widgets and extract their _canvas_manager reference.
    // CanvasManagerModel is a headless singleton that isn't added to the store,
    // but its ID is referenced by each CanvasModel via _canvas_manager.
    models.forEach((model, _id) => {
      if (model.modelName !== "CanvasModel") return;

      const managerRef = model.state?._canvas_manager as string | undefined;
      if (!managerRef) return;

      // Extract manager ID from "IPY_MODEL_xxx" reference
      const managerId = managerRef.startsWith("IPY_MODEL_")
        ? managerRef.slice(10)
        : managerRef;

      if (activeRoutes.has(managerId)) return;

      activeRoutes.set(managerId, setupManagerRouting(store, managerId));
    });

    // Clean up routes for managers no longer referenced by any canvas.
    const referencedManagers = new Set<string>();
    models.forEach((model) => {
      if (model.modelName === "CanvasModel") {
        const managerRef = model.state?._canvas_manager as string | undefined;
        if (managerRef) {
          const managerId = managerRef.startsWith("IPY_MODEL_")
            ? managerRef.slice(10)
            : managerRef;
          referencedManagers.add(managerId);
        }
      }
    });

    for (const [managerId, cleanup] of activeRoutes) {
      if (!referencedManagers.has(managerId)) {
        cleanup();
        activeRoutes.delete(managerId);
      }
    }
  }

  const unsubscribe = store.subscribe(scan);
  scan();

  return () => {
    unsubscribe();
    activeRoutes.forEach((cleanup) => cleanup());
    activeRoutes.clear();
  };
}