CedarCopilot/cedar-mail

import { useMemo } from 'react';
import { useQuery } from '@tanstack/react-query';
import { useTRPC } from '@/providers/query-provider';
import { useSession } from '@/modules/auth/utils/auth-client';
import { getInboxSlug, useInboxes } from '@/modules/threads/hooks/use-inboxes';
import { useInboxCompiledQueries } from '@/modules/threads/hooks/use-inbox-compiled-queries';
import { unreadOnlyQueryHash, withUnreadOnly } from '@/modules/threads/lib/unread-filter';
import { useCedarStore } from '@/modules/store';

/**
 * Per-inbox thread counts for the tab badges sourced from Gmail query counts.
 * `byId` is keyed by inbox id (including system inboxes), and `done` is the
 * archived-thread estimate from Gmail.
 */
export function useInboxCounts() {
  const trpc = useTRPC();
  const { data: session } = useSession();
  const { inboxes } = useInboxes();
  const compiledByInboxId = useInboxCompiledQueries();
  // ⇧U-filtered tabs count their UNREAD, not their all-mail total: the badge names the list
  // the tab opens, and that tab now opens a filtered one.
  const unreadOnlyByFolder = useCedarStore((state) => state.unreadOnlyByFolder);

  const splitInputs = useMemo(
    () =>
      inboxes.map((inbox) => {
        const { compiledQuery, queryHash, inboxName } = compiledByInboxId[inbox.id]!;
        const unreadOnly = !!unreadOnlyByFolder[getInboxSlug(inbox)];
        // A CRM-filtered inbox cannot be counted through Gmail — the server looks its rule up by
        // `inboxId` and routes it to an exact SQL count over the mirror instead, which is what
        // makes the badge agree with the list the tab opens.
        return {
          inboxId: inbox.id,
          inboxName,
          compiledQuery: withUnreadOnly(compiledQuery, unreadOnly),
          queryHash: unreadOnlyQueryHash(queryHash, unreadOnly),
        };
      }),
    [inboxes, compiledByInboxId, unreadOnlyByFolder],
  );

  const query = useQuery(
    trpc.mail.getSplitCounts.queryOptions(
      { splits: splitInputs },
      {
        enabled: !!session?.user.id,
        staleTime: 30_000,
        // ⇧U re-keys this query too. Holding the previous counts keeps the badges in place
        // for the one round trip rather than blanking every tab, then they settle on the
        // filtered numbers.
        placeholderData: (previousData) => previousData,
      },
    ),
  );

  return useMemo(() => {
    const byId: Record<string, { count: number; isExact: boolean }> = { ...query.data?.byId };
    const done = query.data?.done ?? { count: 0, isExact: true };
    return { byId, done, isLoading: query.isLoading };
  }, [query.data, query.isLoading]);
}