cinny/src/app/features/message-search/useMessageSearch.ts

import {
  IEventWithRoomId,
  IResultContext,
  ISearchRequestBody,
  ISearchResponse,
  ISearchResult,
  MatrixClient,
  MatrixEvent,
  Room,
  SearchOrderBy,
} from 'matrix-js-sdk';
import { useCallback } from 'react';
import { useMatrixClient } from '../../hooks/useMatrixClient';

export type ResultItem = {
  rank: number;
  event: IEventWithRoomId;
  context: IResultContext;
};

export type ResultGroup = {
  roomId: string;
  items: ResultItem[];
};

export type SearchResult = {
  nextToken?: string;
  highlights: string[];
  groups: ResultGroup[];
};

const EMPTY_CONTEXT: IResultContext = {
  events_before: [],
  events_after: [],
  profile_info: {},
};

const groupSearchResult = (results: ISearchResult[]): ResultGroup[] => {
  const groups: ResultGroup[] = [];

  results.forEach((item) => {
    const roomId = item.result.room_id;
    const resultItem: ResultItem = {
      rank: item.rank,
      event: item.result,
      context: item.context,
    };

    const lastAddedGroup: ResultGroup | undefined = groups[groups.length - 1];
    if (lastAddedGroup && roomId === lastAddedGroup.roomId) {
      lastAddedGroup.items.push(resultItem);
      return;
    }
    groups.push({
      roomId,
      items: [resultItem],
    });
  });

  return groups;
};

const parseSearchResult = (result: ISearchResponse): SearchResult => {
  const roomEvents = result.search_categories.room_events;

  return {
    nextToken: roomEvents?.next_batch,
    highlights: roomEvents?.highlights ?? [],
    groups: groupSearchResult(roomEvents?.results ?? []),
  };
};

const eventToSearchEvent = (event: MatrixEvent, roomId: string): IEventWithRoomId | undefined => {
  const eventId = event.getId();
  if (!eventId) return undefined;

  const content = event.getClearContent() ?? event.getContent();
  return {
    event_id: eventId,
    type: event.getWireType() === 'm.room.encrypted' ? 'm.room.message' : event.getType(),
    sender: event.getSender() ?? '',
    origin_server_ts: event.getTs(),
    content,
    room_id: roomId,
    unsigned: event.getUnsigned(),
  };
};

const getSearchableBody = (event: MatrixEvent): string | undefined => {
  if (event.isRedacted()) return undefined;

  // After decryption, clear content is available even if wire type was encrypted
  const content = event.getClearContent() ?? event.getContent();
  if (!content || typeof content !== 'object') return undefined;

  // Skip still-encrypted payloads
  if (event.isEncrypted() && !event.isDecryptionFailure() && !event.getClearContent()) {
    return undefined;
  }

  const msgType = content.msgtype;
  if (msgType && msgType !== 'm.text' && msgType !== 'm.notice' && msgType !== 'm.emote') {
    // Still allow filename / body on media
  }

  const body = typeof content.body === 'string' ? content.body : undefined;
  const formatted =
    typeof content.formatted_body === 'string' ? content.formatted_body : undefined;
  return [body, formatted].filter(Boolean).join('\n') || undefined;
};

const highlightsFromTerm = (term: string): string[] =>
  term
    .split(/\s+/)
    .map((part) => part.trim())
    .filter((part) => part.length > 1);

/**
 * Search decrypted timeline events already loaded for a room.
 * Server-side search cannot match E2EE message bodies.
 */
const searchLocalRoomTimeline = (room: Room, term: string): ResultItem[] => {
  const needle = term.trim().toLowerCase();
  if (!needle) return [];

  const events = room.getLiveTimeline().getEvents();
  const matches: ResultItem[] = [];

  for (let i = events.length - 1; i >= 0; i -= 1) {
    const event = events[i];
    const haystack = getSearchableBody(event);
    if (!haystack || !haystack.toLowerCase().includes(needle)) continue;

    const searchEvent = eventToSearchEvent(event, room.roomId);
    if (!searchEvent) continue;

    matches.push({
      rank: 1,
      event: searchEvent,
      context: EMPTY_CONTEXT,
    });
  }

  return matches;
};

const isEncryptedRoom = (mx: MatrixClient, roomId: string): boolean => {
  const room = mx.getRoom(roomId);
  return !!room?.hasEncryptionStateEvent();
};

const searchLocalRooms = (mx: MatrixClient, roomIds: string[], term: string): SearchResult => {
  const groups: ResultGroup[] = [];

  roomIds.forEach((roomId) => {
    const room = mx.getRoom(roomId);
    if (!room) return;
    const items = searchLocalRoomTimeline(room, term);
    if (items.length > 0) {
      groups.push({ roomId, items });
    }
  });

  return {
    highlights: highlightsFromTerm(term),
    groups,
  };
};

export type MessageSearchParams = {
  term?: string;
  order?: string;
  rooms?: string[];
  senders?: string[];
};

export const useMessageSearch = (params: MessageSearchParams) => {
  const mx = useMatrixClient();
  const { term, order, rooms, senders } = params;

  const searchMessages = useCallback(
    async (nextBatch?: string) => {
      if (!term) {
        return {
          highlights: [],
          groups: [],
        };
      }

      const scopedRooms = rooms?.filter(Boolean) ?? [];
      const encryptedScoped =
        scopedRooms.length > 0 && scopedRooms.every((roomId) => isEncryptedRoom(mx, roomId));

      // Encrypted-only scope: server cannot see bodies — search loaded timeline locally
      if (encryptedScoped && !nextBatch) {
        return searchLocalRooms(mx, scopedRooms, term);
      }

      const limit = 20;
      const requestBody: ISearchRequestBody = {
        search_categories: {
          room_events: {
            event_context: {
              before_limit: 0,
              after_limit: 0,
              include_profile: false,
            },
            filter: {
              limit,
              rooms,
              senders,
            },
            include_state: false,
            order_by: order as SearchOrderBy.Recent,
            search_term: term,
          },
        },
      };

      try {
        const r = await mx.search({
          body: requestBody,
          next_batch: nextBatch === '' ? undefined : nextBatch,
        });
        const parsed = parseSearchResult(r);

        if (
          parsed.groups.length === 0 &&
          !nextBatch &&
          scopedRooms.length > 0 &&
          scopedRooms.some((roomId) => isEncryptedRoom(mx, roomId))
        ) {
          return searchLocalRooms(
            mx,
            scopedRooms.filter((id) => isEncryptedRoom(mx, id)),
            term
          );
        }

        return parsed;
      } catch {
        if (!nextBatch && scopedRooms.length > 0) {
          return searchLocalRooms(mx, scopedRooms, term);
        }
        throw new Error('Message search failed');
      }
    },
    [mx, term, order, rooms, senders]
  );

  return searchMessages;
};