mnote/wolai-frontend/src/components/sidebar/use-preferred-sidebar-snapshot.ts

import { useMemo } from "react";
import type { SidebarInitialData } from "@/components/sidebar/types";
import { buildSidebarDataSyncKey } from "@/components/sidebar/sidebar-sync";

export type PreferredSidebarSnapshotSource = "initial" | "query" | "tree_stream";

export function usePreferredSidebarSnapshot(input: {
  initialData: SidebarInitialData;
  sidebarQueryData: SidebarInitialData | null;
  treeStreamData: SidebarInitialData | null;
  treeStreamStatus?: "idle" | "connecting" | "live" | "fallback";
}) {
  const querySyncKey = useMemo(
    () => (input.sidebarQueryData ? buildSidebarDataSyncKey(input.sidebarQueryData) : null),
    [input.sidebarQueryData],
  );
  const treeStreamSyncKey = useMemo(
    () => (input.treeStreamData ? buildSidebarDataSyncKey(input.treeStreamData) : null),
    [input.treeStreamData],
  );
  const initialSyncKey = useMemo(() => buildSidebarDataSyncKey(input.initialData), [input.initialData]);
  const streamIsPreferred = input.treeStreamStatus !== "fallback";

  const source = useMemo<PreferredSidebarSnapshotSource>(() => {
    if (input.treeStreamData && streamIsPreferred) {
      return "tree_stream";
    }
    if (input.sidebarQueryData) {
      return "query";
    }
    return "initial";
  }, [
    input.sidebarQueryData,
    input.treeStreamData,
    streamIsPreferred,
  ]);

  const data =
    source === "tree_stream" && input.treeStreamData
      ? input.treeStreamData
      : source === "query" && input.sidebarQueryData
        ? input.sidebarQueryData
        : input.treeStreamData ?? input.sidebarQueryData ?? input.initialData;
  const syncKey =
    source === "tree_stream"
      ? treeStreamSyncKey ?? querySyncKey ?? initialSyncKey
      : source === "query"
        ? querySyncKey ?? initialSyncKey
        : initialSyncKey;

  return useMemo(
    () => ({
      data,
      source,
      syncKey,
    }),
    [data, source, syncKey],
  );
}