yaak/src-web/hooks/useKeyValue.ts

import deepEqual from '@gilbarbara/deep-equal';
import { useMutation } from '@tanstack/react-query';
import { keyValuesAtom } from '@yaakapp-internal/models';
import { useAtomValue } from 'jotai';
import { selectAtom } from 'jotai/utils';
import { useCallback, useMemo } from 'react';
import { jotaiStore } from '../lib/jotai';
import { buildKeyValueKey, extractKeyValueOrFallback, setKeyValue } from '../lib/keyValueStore';

const DEFAULT_NAMESPACE = 'global';

export function useKeyValue<T extends object | boolean | number | string | null>({
  namespace = DEFAULT_NAMESPACE,
  key,
  fallback,
}: {
  namespace?: 'global' | 'no_sync' | 'license';
  key: string | string[];
  fallback: T;
}) {
  const { value, isLoading } = useAtomValue(
    // biome-ignore lint/correctness/useExhaustiveDependencies: Only create a new atom when the key changes. Fallback might not be a stable reference, so we don't want to refresh on that.
    useMemo(
      () =>
        selectAtom(
          keyValuesAtom,
          (keyValues) => {
            const keyValue =
              keyValues?.find((kv) => buildKeyValueKey(kv.key) === buildKeyValueKey(key)) ?? null;
            const value = keyValues == null ? null : extractKeyValueOrFallback(keyValue, fallback);
            const isLoading = keyValues == null;
            return { value, isLoading };
          },
          (a, b) => deepEqual(a, b),
        ),
      [buildKeyValueKey(key)],
    ),
  );

  const { mutateAsync } = useMutation<void, unknown, T>({
    mutationKey: ['set_key_value', namespace, key],
    mutationFn: (value) => setKeyValue<T>({ namespace, key, value }),
  });

  // biome-ignore lint/correctness/useExhaustiveDependencies: none
  const set = useCallback(
    async (valueOrUpdate: ((v: T) => T) | T) => {
      if (typeof valueOrUpdate === 'function') {
        const newV = valueOrUpdate(value ?? fallback);
        if (newV === value) return;
        await mutateAsync(newV);
      } else {
        // TODO: Make this only update if the value is different. I tried this but it seems query.data
        //  is stale.
        await mutateAsync(valueOrUpdate);
      }
    },
    [typeof key === 'string' ? key : key.join('::'), namespace, value],
  );

  const reset = useCallback(async () => mutateAsync(fallback), [fallback, mutateAsync]);

  return useMemo(
    () => ({
      value,
      isLoading,
      set,
      reset,
    }),
    [isLoading, reset, set, value],
  );
}

export function getKeyValue<T extends object | boolean | number | string | null>({
  namespace,
  key,
  fallback,
}: {
  namespace?: 'global' | 'no_sync' | 'license';
  key: string | string[];
  fallback: T;
}) {
  const keyValues = jotaiStore.get(keyValuesAtom);
  const keyValue =
    keyValues?.find(
      (kv) => kv.namespace === namespace && buildKeyValueKey(kv.key) === buildKeyValueKey(key),
    ) ?? null;
  const value = extractKeyValueOrFallback(keyValue, fallback);
  return value;
}