memos/web/src/hooks/useMemoQueries.ts

import { create } from "@bufbuild/protobuf";
import { FieldMaskSchema } from "@bufbuild/protobuf/wkt";
import { Code, ConnectError } from "@connectrpc/connect";
import {
  type InfiniteData,
  type QueryClient,
  queryOptions,
  useInfiniteQuery,
  useMutation,
  useQuery,
  useQueryClient,
} from "@tanstack/react-query";
import { memoServiceClient } from "@/connect";
import { attachmentKeys } from "@/hooks/useAttachmentQueries";
import { userKeys } from "@/hooks/useUserQueries";
import { DEFAULT_LIST_MEMOS_PAGE_SIZE } from "@/lib/constants";
import { shouldRetry } from "@/lib/query-client";
import type { ListMemosRequest, ListMemosResponse, Memo } from "@/types/proto/api/v1/memo_service_pb";
import { ListMemoCommentsRequestSchema, ListMemosRequestSchema, MemoSchema } from "@/types/proto/api/v1/memo_service_pb";

// Query keys factory for consistent cache management
export const memoKeys = {
  all: ["memos"] as const,
  lists: () => [...memoKeys.all, "list"] as const,
  list: (filters: Partial<ListMemosRequest>) => [...memoKeys.lists(), filters] as const,
  details: () => [...memoKeys.all, "detail"] as const,
  detail: (name: string) => [...memoKeys.details(), name] as const,
  comments: (name: string) => [...memoKeys.all, "comments", name] as const,
  linkMetadata: (url: string) => [...memoKeys.all, "linkMetadata", url] as const,
};

export const memoDetailQueryOptions = (name: string) =>
  queryOptions({
    queryKey: memoKeys.detail(name),
    queryFn: async ({ client, signal }) => {
      try {
        return await memoServiceClient.getMemo({ name }, { signal });
      } catch (error) {
        if (!isMemoUnavailableError(error)) throw error;
        // Store a content-free result, replacing any formerly readable detail.
        // Remove copies and snippets so collection fallbacks cannot revive them.
        discardUnavailableMemo(client, name);
        return null;
      }
    },
    retry: shouldRetry,
    staleTime: 1000 * 10,
    // A previous denial cannot establish access on a new visit, even while fresh.
    refetchOnMount: (query) => (query.state.data === null ? "always" : true),
  });

export function isMemoUnavailableError(error: unknown): boolean {
  return error instanceof ConnectError && (error.code === Code.PermissionDenied || error.code === Code.NotFound);
}

function discardUnavailableMemo(client: QueryClient, name: string) {
  const stripRelations = (memo: Memo): Memo => {
    const relations = memo.relations.filter((relation) => relation.memo?.name !== name && relation.relatedMemo?.name !== name);
    return relations.length === memo.relations.length ? memo : { ...memo, relations };
  };
  const stripList = (data: ListMemosResponse): ListMemosResponse => {
    const memos = data.memos.filter((memo) => memo.name !== name).map(stripRelations);
    return memos.length === data.memos.length && memos.every((memo, index) => memo === data.memos[index]) ? data : { ...data, memos };
  };
  for (const [key, data] of client.getQueriesData<MemoCollectionQueryData | Memo | null>({ queryKey: memoKeys.all })) {
    let next = data;
    if (isMemoListResponse(data)) next = stripList(data);
    else if (isInfiniteMemoListData(data)) {
      const pages = data.pages.map(stripList);
      if (pages.some((page, index) => page !== data.pages[index])) next = { ...data, pages };
    } else if (key[1] === "detail" && data && "name" in data && data.name !== name) next = stripRelations(data);
    if (next !== data) client.setQueryData(key, next);
  }
}

type MemoPatch = Partial<Memo> & Pick<Memo, "name">;
type MemoCollectionQueryData = ListMemosResponse | InfiniteData<ListMemosResponse>;

function isMemoListResponse(data: unknown): data is ListMemosResponse {
  return typeof data === "object" && data !== null && Array.isArray((data as { memos?: unknown }).memos);
}

function isInfiniteMemoListData(data: unknown): data is InfiniteData<ListMemosResponse> {
  return typeof data === "object" && data !== null && Array.isArray((data as { pages?: unknown }).pages);
}

function patchMemoListResponse(response: ListMemosResponse, update: MemoPatch): ListMemosResponse {
  let changed = false;
  const memos = response.memos.map((memo) => {
    if (memo.name !== update.name) {
      return memo;
    }

    changed = true;
    return { ...memo, ...update };
  });

  return changed ? { ...response, memos } : response;
}

function patchMemoListQueryData<T>(data: T | undefined, update: MemoPatch): T | undefined {
  if (!data) {
    return data;
  }

  if (isMemoListResponse(data)) {
    return patchMemoListResponse(data, update) as T;
  }

  if (isInfiniteMemoListData(data)) {
    let changed = false;
    const pages = data.pages.map((page) => {
      const patchedPage = patchMemoListResponse(page, update);
      if (patchedPage !== page) {
        changed = true;
      }
      return patchedPage;
    });

    return (changed ? { ...data, pages } : data) as T;
  }

  return data;
}

function findMemoInListResponse(response: ListMemosResponse, name: string): Memo | undefined {
  return response.memos.find((memo) => memo.name === name);
}

function findMemoInQueryData(data: unknown, name: string): Memo | undefined {
  if (!data) {
    return undefined;
  }

  if (isMemoListResponse(data)) {
    return findMemoInListResponse(data, name);
  }

  if (isInfiniteMemoListData(data)) {
    for (const page of data.pages) {
      const memo = findMemoInListResponse(page, name);
      if (memo) {
        return memo;
      }
    }
  }

  return undefined;
}

export function findMemoInCollectionQueries(queryClient: QueryClient, name: string, freshOnly = false): Memo | undefined {
  if (queryClient.getQueryData(memoKeys.detail(name)) === null) return undefined;
  for (const [key, data] of queryClient.getQueriesData<unknown>({ queryKey: memoKeys.all })) {
    const state = queryClient.getQueryState(key);
    if (freshOnly && (!state || state.isInvalidated || state.dataUpdatedAt < Date.now() - 10_000)) continue;
    const memo = findMemoInQueryData(data, name);
    if (memo) {
      return memo;
    }
  }

  return undefined;
}

function patchMemoInCollectionQueries(queryClient: QueryClient, update: MemoPatch) {
  queryClient.setQueriesData<MemoCollectionQueryData>({ queryKey: memoKeys.all }, (data) => patchMemoListQueryData(data, update));
}

export function useMemos(request: Partial<ListMemosRequest> = {}) {
  return useQuery({
    queryKey: memoKeys.list(request),
    queryFn: async () => {
      const response = await memoServiceClient.listMemos(create(ListMemosRequestSchema, request as Record<string, unknown>));
      return response;
    },
  });
}

export function useInfiniteMemos(request: Partial<ListMemosRequest> = {}, options?: { enabled?: boolean }) {
  return useInfiniteQuery({
    queryKey: memoKeys.list(request),
    queryFn: async ({ pageParam }) => {
      const response = await memoServiceClient.listMemos(
        create(ListMemosRequestSchema, {
          ...request,
          pageToken: pageParam || "",
        } as Record<string, unknown>),
      );
      return response;
    },
    initialPageParam: "",
    getNextPageParam: (lastPage) => lastPage.nextPageToken || undefined,
    staleTime: 1000 * 60,
    gcTime: 1000 * 60 * 5,
    enabled: options?.enabled ?? true,
  });
}

export function useMemo(name: string, options?: { enabled?: boolean }) {
  const query = useQuery({
    ...memoDetailQueryOptions(name),
    enabled: options?.enabled ?? true,
  });
  return { ...query, data: query.data ?? undefined, isUnavailable: query.data === null };
}

function isHTTPURL(url: string): boolean {
  try {
    const parsed = new URL(url);
    return parsed.protocol === "http:" || parsed.protocol === "https:";
  } catch {
    return false;
  }
}

export function useLinkMetadata(url: string, options?: { enabled?: boolean }) {
  const trimmedUrl = url.trim();

  return useQuery({
    queryKey: memoKeys.linkMetadata(trimmedUrl),
    queryFn: async () => {
      const metadata = await memoServiceClient.getLinkMetadata({ url: trimmedUrl });
      return metadata;
    },
    enabled: (options?.enabled ?? true) && isHTTPURL(trimmedUrl),
    staleTime: 1000 * 60 * 60 * 24,
    gcTime: 1000 * 60 * 60 * 24,
  });
}

export function useCreateMemo() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: async (memoToCreate: Memo) => {
      const memo = await memoServiceClient.createMemo({ memo: memoToCreate });
      return memo;
    },
    onSuccess: (newMemo) => {
      // Invalidate memo lists to refetch
      queryClient.invalidateQueries({ queryKey: memoKeys.lists() });
      // Add new memo to cache
      queryClient.setQueryData(memoKeys.detail(newMemo.name), newMemo);
      // Invalidate user stats
      queryClient.invalidateQueries({ queryKey: userKeys.stats() });
      // Creating a memo can bind previously unlinked attachments.
      queryClient.invalidateQueries({ queryKey: attachmentKeys.lists() });
    },
  });
}

export function useUpdateMemo() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: async ({ update, updateMask }: { update: Partial<Memo>; updateMask: string[] }) => {
      const memo = await memoServiceClient.updateMemo({
        memo: create(MemoSchema, update as Record<string, unknown>),
        updateMask: create(FieldMaskSchema, { paths: updateMask }),
      });
      return memo;
    },
    onMutate: async ({ update, updateMask }) => {
      if (updateMask.includes("space")) return { previousMemo: undefined };
      if (!update.name) {
        return { previousMemo: undefined };
      }

      // Cancel outgoing refetches to prevent race conditions
      await queryClient.cancelQueries({ queryKey: memoKeys.all });

      // Snapshot previous value for rollback on error
      const previousMemo =
        queryClient.getQueryData<Memo>(memoKeys.detail(update.name)) || findMemoInCollectionQueries(queryClient, update.name);
      const memoPatch: MemoPatch = { ...update, name: update.name };

      // Optimistically update the cache
      if (previousMemo) {
        queryClient.setQueryData(memoKeys.detail(update.name), { ...previousMemo, ...memoPatch });
      }
      patchMemoInCollectionQueries(queryClient, memoPatch);

      return { previousMemo };
    },
    onError: (_err, { update }, context) => {
      // Rollback on error
      if (context?.previousMemo && update.name) {
        queryClient.setQueryData(memoKeys.detail(update.name), context.previousMemo);
        patchMemoInCollectionQueries(queryClient, context.previousMemo);
      } else {
        queryClient.invalidateQueries({ queryKey: memoKeys.all });
      }
    },
    onSuccess: (updatedMemo, { updateMask }) => {
      if (updateMask.includes("space")) {
        queryClient.invalidateQueries({ queryKey: memoKeys.all });
      }
      // Update cache with server response
      queryClient.setQueryData(memoKeys.detail(updatedMemo.name), updatedMemo);
      patchMemoInCollectionQueries(queryClient, updatedMemo);
      // Invalidate lists to refresh
      queryClient.invalidateQueries({ queryKey: memoKeys.lists() });
      if (updatedMemo.parent) {
        queryClient.invalidateQueries({ queryKey: memoKeys.comments(updatedMemo.parent) });
      }
      // Invalidate user stats
      queryClient.invalidateQueries({ queryKey: userKeys.stats() });
      // Placement changes move linked attachments between scoped libraries.
      queryClient.invalidateQueries({ queryKey: attachmentKeys.lists() });
    },
  });
}

export function useDeleteMemo() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: async (name: string) => {
      await memoServiceClient.deleteMemo({ name });
      return name;
    },
    onSuccess: (name) => {
      // Remove from cache
      queryClient.removeQueries({ queryKey: memoKeys.detail(name) });
      // Invalidate lists
      queryClient.invalidateQueries({ queryKey: memoKeys.lists() });
      // Invalidate user stats
      queryClient.invalidateQueries({ queryKey: userKeys.stats() });
      // Memo deletion can remove or unlink associated attachments.
      queryClient.invalidateQueries({ queryKey: attachmentKeys.lists() });
    },
  });
}

export function useMemoComments(name: string, options?: { enabled?: boolean; pageSize?: number }) {
  return useQuery({
    queryKey: [...memoKeys.comments(name), options?.pageSize ?? 0],
    queryFn: async () => {
      const response = await memoServiceClient.listMemoComments(
        create(ListMemoCommentsRequestSchema, {
          name,
          pageSize: options?.pageSize ?? 0,
        }),
      );
      return response;
    },
    enabled: options?.enabled ?? true,
    staleTime: 1000 * 60, // 1 minute
  });
}

// useInfiniteMemoComments paginates through every comment via nextPageToken, instead of
// stopping at the server's default page size (the cause of comments being truncated to 10).
export function useInfiniteMemoComments(name: string, options?: { enabled?: boolean; pageSize?: number }) {
  const pageSize = options?.pageSize ?? DEFAULT_LIST_MEMOS_PAGE_SIZE;
  return useInfiniteQuery({
    queryKey: [...memoKeys.comments(name), "infinite", pageSize],
    queryFn: async ({ pageParam }) => {
      const response = await memoServiceClient.listMemoComments(
        create(ListMemoCommentsRequestSchema, {
          name,
          pageSize,
          pageToken: pageParam || "",
        }),
      );
      return response;
    },
    initialPageParam: "",
    getNextPageParam: (lastPage) => lastPage.nextPageToken || undefined,
    select: (data) => data.pages.flatMap((page) => page.memos),
    enabled: options?.enabled ?? true,
    staleTime: 1000 * 60, // 1 minute
  });
}