cinny/src/app/hooks/useThreadNotifications.ts

import { useCallback } from 'react';
import { useAtomValue } from 'jotai';
import { MatrixClient } from 'matrix-js-sdk';
import { AccountDataEvent } from '../../types/matrix/accountData';
import { threadNotificationsAtom } from '../state/threadNotifications';
import {
  getThreadNotificationMode,
  pruneThreadNotifications,
  ThreadNotificationEntry,
  ThreadNotificationMode,
  ThreadNotificationsContent,
} from '../utils/threadNotifications';
import { useMatrixClient } from './useMatrixClient';
import { AsyncState, useAsyncCallback } from './useAsyncCallback';
import { getAccountData, setAccountData } from '../utils/accountData';

/** Read the current notification mode for a thread from the bound atom. */
export function useThreadNotificationMode(
  roomId: string,
  threadRootId: string,
): ThreadNotificationMode {
  const content = useAtomValue(threadNotificationsAtom);
  return getThreadNotificationMode(content, roomId, threadRootId);
}

const readContent = (mx: MatrixClient): ThreadNotificationsContent =>
  getAccountData<ThreadNotificationsContent>(mx, AccountDataEvent.LotusThreadNotifications) ?? {};

const getJoinedRoomIds = (mx: MatrixClient): Set<string> => {
  const joined = new Set<string>();
  mx.getRooms().forEach((room) => {
    if (room.getMyMembership() === 'join') {
      joined.add(room.roomId);
    }
  });
  return joined;
};

// Apply a single mode change to a base content object, returning a fresh clone
// (never mutates the input).
const applyThreadMode = (
  base: ThreadNotificationsContent,
  roomId: string,
  threadRootId: string,
  mode: ThreadNotificationMode,
  now: number,
): ThreadNotificationsContent => {
  const next: ThreadNotificationsContent = {
    ...base,
    rooms: Object.fromEntries(
      Object.entries(base.rooms ?? {}).map(([rid, entries]) => [rid, { ...entries }]),
    ),
  };

  const rooms = next.rooms as Record<string, Record<string, ThreadNotificationEntry>>;

  if (mode === ThreadNotificationMode.Default) {
    if (rooms[roomId]) {
      delete rooms[roomId][threadRootId];
      if (Object.keys(rooms[roomId]).length === 0) {
        delete rooms[roomId];
      }
    }
  } else {
    if (!rooms[roomId]) {
      rooms[roomId] = {};
    }
    rooms[roomId][threadRootId] = { mode, ts: now };
  }

  return next;
};

// T7 — serialize writes so rapid, overlapping mode changes don't lost-update
// each other. `setAccountData` is a bare PUT whose result doesn't reach the
// local store until the /sync echo, so back-to-back writes would otherwise all
// read the same stale base and clobber one another. Each queued write instead
// bases its mutation on the previous write's RESULT; once the queue drains the
// carried base is dropped so the next independent write re-reads fresh (possibly
// externally-changed) server state.
let writeChain: Promise<unknown> = Promise.resolve();
let pendingWrites = 0;
let carriedContent: ThreadNotificationsContent | null = null;

const writeThreadNotificationMode = (
  mx: MatrixClient,
  roomId: string,
  threadRootId: string,
  mode: ThreadNotificationMode,
): Promise<void> => {
  pendingWrites += 1;
  const run = writeChain.then(async () => {
    const now = Date.now();
    const base = carriedContent ?? readContent(mx);
    const next = applyThreadMode(base, roomId, threadRootId, mode, now);
    // ALWAYS prune before persisting to keep account data bounded.
    const finalContent = pruneThreadNotifications(next, getJoinedRoomIds(mx), now);
    await setAccountData(mx, AccountDataEvent.LotusThreadNotifications, finalContent);
    // Carry the result forward only on success, so a queued follow-up write bases
    // on persisted content — never on a shape the server just rejected.
    carriedContent = finalContent;
  });
  // Keep the chain alive on error; drop the carried base once the queue drains.
  writeChain = run
    .catch(() => {})
    .finally(() => {
      pendingWrites -= 1;
      if (pendingWrites === 0) carriedContent = null;
    });
  return run;
};

export function useSetThreadNotificationMode(
  roomId: string,
  threadRootId: string,
): {
  modeState: AsyncState<void, Error>;
  setMode: (mode: ThreadNotificationMode) => Promise<void>;
} {
  const mx = useMatrixClient();

  const [modeState, setMode] = useAsyncCallback<void, Error, [ThreadNotificationMode]>(
    useCallback(
      (mode: ThreadNotificationMode) => writeThreadNotificationMode(mx, roomId, threadRootId, mode),
      [mx, roomId, threadRootId],
    ),
  );

  return { modeState, setMode };
}