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;
}) {
  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 queryFreshness = useMemo(
    () => (input.sidebarQueryData ? getSidebarDataFreshness(input.sidebarQueryData) : Number.NEGATIVE_INFINITY),
    [input.sidebarQueryData],
  );
  const treeStreamFreshness = useMemo(
    () => (input.treeStreamData ? getSidebarDataFreshness(input.treeStreamData) : Number.NEGATIVE_INFINITY),
    [input.treeStreamData],
  );

  const source = useMemo<PreferredSidebarSnapshotSource>(() => {
    if (input.treeStreamData && input.sidebarQueryData) {
      if (treeStreamSyncKey === querySyncKey) {
        return "tree_stream";
      }
      return queryFreshness > treeStreamFreshness ? "query" : "tree_stream";
    }
    if (input.treeStreamData) {
      return "tree_stream";
    }
    if (input.sidebarQueryData) {
      return "query";
    }
    return "initial";
  }, [
    input.sidebarQueryData,
    input.treeStreamData,
    queryFreshness,
    querySyncKey,
    treeStreamFreshness,
    treeStreamSyncKey,
  ]);

  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 ?? initialSyncKey
      : source === "query"
        ? querySyncKey ?? initialSyncKey
        : initialSyncKey;

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