cinny/src/app/hooks/useRoomMediaEvents.ts

import { useCallback, useEffect, useRef, useState } from 'react';
import to from 'await-to-js';
import {
  Direction,
  IRoomEvent,
  MatrixEvent,
  MsgType,
  RelationType,
  Room,
} from 'matrix-js-sdk';
import { CryptoBackend } from 'matrix-js-sdk/lib/common-crypto/CryptoBackend';
import { EncryptedAttachmentInfo } from 'browser-encrypt-attachment';
import { useMatrixClient } from './useMatrixClient';
import { IEncryptedFile } from '../../types/matrix/common';

export type RoomMediaItem = {
  eventId: string;
  msgtype: typeof MsgType.Image | typeof MsgType.Video;
  /** MXC used for an image thumbnail when available. */
  thumbMxc?: string;
  thumbMimeType?: string;
  thumbEncInfo?: EncryptedAttachmentInfo;
  /** Full media MXC — used to extract a video frame when no image thumb exists. */
  fullMxc: string;
  fullMimeType?: string;
  fullEncInfo?: EncryptedAttachmentInfo;
  body: string;
  ts: number;
  sender: string;
};

const PAGE_SIZE = 50;
const TARGET_BATCH = 24;
const MAX_PAGES_PER_LOAD = 8;

type RoomMediaCache = {
  items: RoomMediaItem[];
  hasMore: boolean;
  nextToken: string | null;
  seenIds: Set<string>;
};

/** Survives MediaDrawer remounts within the same room/session. */
const mediaCache = new Map<string, RoomMediaCache>();

const getOrCreateCache = (roomId: string): RoomMediaCache => {
  const existing = mediaCache.get(roomId);
  if (existing) return existing;
  const created: RoomMediaCache = {
    items: [],
    hasMore: true,
    nextToken: null,
    seenIds: new Set(),
  };
  mediaCache.set(roomId, created);
  return created;
};

const toMediaItem = (mEvent: MatrixEvent): RoomMediaItem | undefined => {
  if (mEvent.isRedacted()) return undefined;

  const relation = mEvent.getRelation();
  if (relation?.rel_type === RelationType.Replace) return undefined;

  const content = mEvent.getContent();
  const msgtype = content.msgtype;
  if (msgtype !== MsgType.Image && msgtype !== MsgType.Video) return undefined;

  const file = content.file as IEncryptedFile | undefined;
  const info = content.info as
    | {
        mimetype?: string;
        thumbnail_url?: string;
        thumbnail_info?: { mimetype?: string };
        thumbnail_file?: IEncryptedFile;
      }
    | undefined;

  const fullMxc =
    (typeof content.url === 'string' && content.url) ||
    (typeof file?.url === 'string' && file.url) ||
    undefined;
  if (!fullMxc) return undefined;

  const thumbFile = info?.thumbnail_file;
  const thumbMxc =
    (typeof info?.thumbnail_url === 'string' && info.thumbnail_url) ||
    (typeof thumbFile?.url === 'string' && thumbFile.url) ||
    undefined;

  const eventId = mEvent.getId();
  const sender = mEvent.getSender();
  if (!eventId || !sender) return undefined;

  const usingThumbEnc = !!thumbMxc && thumbMxc === thumbFile?.url;

  return {
    eventId,
    msgtype,
    thumbMxc,
    thumbMimeType: usingThumbEnc
      ? info?.thumbnail_info?.mimetype
      : thumbMxc
        ? info?.thumbnail_info?.mimetype
        : undefined,
    thumbEncInfo: usingThumbEnc ? thumbFile : undefined,
    fullMxc,
    fullMimeType: info?.mimetype,
    fullEncInfo: file,
    body: typeof content.body === 'string' ? content.body : msgtype,
    ts: mEvent.getTs(),
    sender,
  };
};

/**
 * Paginates room history for image/video messages without mutating the live timeline.
 */
export const useRoomMediaEvents = (room: Room) => {
  const mx = useMatrixClient();
  const roomId = room.roomId;
  const cacheRef = useRef(getOrCreateCache(roomId));
  if (cacheRef.current !== mediaCache.get(roomId)) {
    cacheRef.current = getOrCreateCache(roomId);
  }

  const [items, setItems] = useState<RoomMediaItem[]>(() => cacheRef.current.items);
  const [loading, setLoading] = useState(false);
  const [hasMore, setHasMore] = useState(() => cacheRef.current.hasMore);
  const loadingRef = useRef(false);

  const decryptEvent = useCallback(
    async (raw: IRoomEvent | MatrixEvent) => {
      const mEvent = raw instanceof MatrixEvent ? raw : new MatrixEvent(raw);
      if (mEvent.isEncrypted() && mx.getCrypto()) {
        await to(mEvent.attemptDecryption(mx.getCrypto() as CryptoBackend));
      }
      return mEvent;
    },
    [mx]
  );

  const fetchPages = useCallback(async () => {
    const cache = cacheRef.current;
    const collected: RoomMediaItem[] = [];

    for (let page = 0; page < MAX_PAGES_PER_LOAD; page += 1) {
      const response = await mx.createMessagesRequest(
        roomId,
        cache.nextToken,
        PAGE_SIZE,
        Direction.Backward
      );

      cache.nextToken = response.end ?? null;
      if (!response.end) {
        cache.hasMore = false;
        setHasMore(false);
      }

      for (const raw of response.chunk ?? []) {
        const mEvent = await decryptEvent(raw);
        const item = toMediaItem(mEvent);
        if (!item || cache.seenIds.has(item.eventId)) continue;
        cache.seenIds.add(item.eventId);
        collected.push(item);
      }

      if (collected.length >= TARGET_BATCH || !response.end) break;
    }

    return collected;
  }, [mx, roomId, decryptEvent]);

  const loadMore = useCallback(async () => {
    const cache = cacheRef.current;
    if (loadingRef.current || !cache.hasMore) return;
    loadingRef.current = true;
    setLoading(true);

    try {
      const collected = await fetchPages();
      if (collected.length > 0) {
        cache.items = [...cache.items, ...collected];
        setItems(cache.items);
      }
    } finally {
      loadingRef.current = false;
      setLoading(false);
    }
  }, [fetchPages]);

  useEffect(() => {
    const cache = getOrCreateCache(roomId);
    cacheRef.current = cache;

    // Already loaded for this room — restore without re-fetching.
    if (cache.items.length > 0 || cache.nextToken !== null || !cache.hasMore) {
      setItems(cache.items);
      setHasMore(cache.hasMore);
      return undefined;
    }

    let cancelled = false;

    const bootstrap = async () => {
      cache.seenIds = new Set();
      cache.nextToken = null;
      cache.hasMore = true;
      cache.items = [];
      setItems([]);
      setHasMore(true);
      setLoading(true);
      loadingRef.current = true;

      try {
        const localFound: RoomMediaItem[] = [];
        const localEvents = room.getLiveTimeline().getEvents();
        for (let i = localEvents.length - 1; i >= 0; i -= 1) {
          const mEvent = await decryptEvent(localEvents[i]);
          const item = toMediaItem(mEvent);
          if (!item || cache.seenIds.has(item.eventId)) continue;
          cache.seenIds.add(item.eventId);
          localFound.push(item);
        }

        if (cancelled) return;
        if (localFound.length > 0) {
          cache.items = localFound;
          setItems(localFound);
        }

        const collected = await fetchPages();
        if (cancelled) return;
        if (collected.length > 0) {
          cache.items = [...cache.items, ...collected];
          setItems(cache.items);
        }
      } finally {
        if (!cancelled) {
          loadingRef.current = false;
          setLoading(false);
        }
      }
    };

    bootstrap();
    return () => {
      cancelled = true;
    };
    // Only re-bootstrap when switching rooms — not when callback identities change.
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [roomId]);

  return { items, loading, hasMore, loadMore };
};