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

import { useMemo } from "react";
import type { SidebarInitialData } from "@/components/sidebar/types";
import { buildSidebarDataSyncKey, getSidebarDataFreshness } 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";
  treeStreamCursor?: string | null;
}) {
  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 initialVersion = useMemo(() => getSidebarDataFreshness(input.initialData), [input.initialData]);
  const queryVersion = useMemo(
    () => (input.sidebarQueryData ? getSidebarDataFreshness(input.sidebarQueryData) : null),
    [input.sidebarQueryData],
  );
  const treeStreamVersion = useMemo(
    () => (input.treeStreamData ? getSidebarDataFreshness(input.treeStreamData) : null),
    [input.treeStreamData],
  );
  const streamIsPreferred = input.treeStreamStatus !== "fallback";

  const preferredVersion = Math.max(
    initialVersion,
    queryVersion ?? 0,
    streamIsPreferred ? treeStreamVersion ?? 0 : 0,
  );
  const queryHasPreferredVersion = queryVersion != null && queryVersion >= preferredVersion;
  const treeStreamHasPreferredVersion =
    streamIsPreferred &&
    input.treeStreamData != null &&
    treeStreamVersion != null &&
    treeStreamVersion >= preferredVersion;

  const source = useMemo<PreferredSidebarSnapshotSource>(() => {
    if (
      queryHasPreferredVersion &&
      (!treeStreamHasPreferredVersion || querySyncKey !== treeStreamSyncKey)
    ) {
      return "query";
    }
    if (treeStreamHasPreferredVersion) {
      return "tree_stream";
    }
    if (queryHasPreferredVersion) {
      return "query";
    }
    return "initial";
  }, [
    queryHasPreferredVersion,
    querySyncKey,
    treeStreamHasPreferredVersion,
    treeStreamSyncKey,
  ]);

  const data =
    source === "tree_stream" && input.treeStreamData
      ? input.treeStreamData
      : source === "query" && input.sidebarQueryData
        ? input.sidebarQueryData
        : input.initialData;
  const syncKey =
    source === "tree_stream"
      ? treeStreamSyncKey ?? querySyncKey ?? initialSyncKey
      : source === "query"
        ? querySyncKey ?? initialSyncKey
        : initialSyncKey;
  const version =
    source === "tree_stream"
      ? treeStreamVersion ?? initialVersion
      : source === "query"
        ? queryVersion ?? initialVersion
        : initialVersion;
  const cursor = source === "tree_stream" ? input.treeStreamCursor ?? null : null;

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