Toju/toju-app/src/app/domains/attachment/application/services/attachment-manager.service.ts

import {
  Injectable,
  effect,
  inject
} from '@angular/core';
import { NavigationEnd, Router } from '@angular/router';
import { Store } from '@ngrx/store';
import { take } from 'rxjs';
import { RealtimeSessionFacade } from '../../../../core/realtime';
import { selectCurrentUserId } from '../../../../store/users/users.selectors';
import { DatabaseService } from '../../../../infrastructure/persistence';
import { yieldToAttachmentHydrationLoop } from '../../domain/logic/attachment-blob.rules';
import {
  buildAttachmentDisplayPinKey,
  collectMessageIdsForInactiveRoomBlobRelease,
  shouldRevokeDisplayBlobForAttachment
} from '../../domain/logic/attachment-blob-eviction.rules';
import {
  getWatchedAttachmentRoomIdFromUrl,
  isDirectMessageAttachmentRoomId,
  shouldAutoRequestWhenWatched
} from '../../domain/logic/attachment.logic';
import { ATTACHMENT_AUTO_DOWNLOAD_CONCURRENCY, runTasksWithBoundedConcurrency } from '../../domain/logic/attachment-autodownload-concurrency.rules';
import { shouldResetStalledAttachmentDownload } from '../../domain/logic/attachment-autodownload.rules';
import type { Attachment, AttachmentMeta } from '../../domain/models/attachment.model';
import type {
  FileAnnouncePayload,
  FileCancelPayload,
  FileChunkPayload,
  FileChunkAckPayload,
  FileNotFoundPayload,
  FileRequestPayload
} from '../../domain/models/attachment-transfer.model';
import { AttachmentPersistenceService } from './attachment-persistence.service';
import { AttachmentRuntimeStore } from './attachment-runtime.store';
import { AttachmentTransferService } from './attachment-transfer.service';

@Injectable({ providedIn: 'root' })
export class AttachmentManagerService {
  get updated() {
    return this.runtimeStore.updated;
  }

  private readonly webrtc = inject(RealtimeSessionFacade);
  private readonly router = inject(Router);
  private readonly store = inject(Store);
  private readonly database = inject(DatabaseService);
  private readonly runtimeStore = inject(AttachmentRuntimeStore);
  private readonly persistence = inject(AttachmentPersistenceService);
  private readonly transfer = inject(AttachmentTransferService);

  private watchedRoomId: string | null = this.extractWatchedRoomId(this.router.url);
  private isDatabaseInitialised = false;
  private autoDownloadRequestsByRoom = new Map<string, Promise<void>>();
  private pinnedDisplayBlobKeys = new Set<string>();

  constructor() {
    effect(() => {
      if (this.database.isReady() && !this.isDatabaseInitialised) {
        this.isDatabaseInitialised = true;
        void this.persistence.initFromDatabase().then(async () => {
          if (this.watchedRoomId) {
            await this.restoreLocalAttachmentsForRoom(this.watchedRoomId);
            await this.announceHostedAttachments();
          }
        });
      }
    });

    this.router.events.subscribe((event) => {
      if (!(event instanceof NavigationEnd)) {
        return;
      }

      const previousRoomId = this.watchedRoomId;

      this.watchedRoomId = this.extractWatchedRoomId(event.urlAfterRedirects || event.url);

      if (this.watchedRoomId !== previousRoomId) {
        this.releaseDisplayBlobsForInactiveRooms(this.watchedRoomId);
      }

      if (this.watchedRoomId) {
        void this.restoreLocalAttachmentsForRoom(this.watchedRoomId);
        void this.requestAutoDownloadsForRoom(this.watchedRoomId);
      }
    });

    this.webrtc.onPeerConnected.subscribe(() => {
      if (this.watchedRoomId) {
        void this.restoreLocalAttachmentsForRoom(this.watchedRoomId).then(async () => {
          await this.announceHostedAttachments();
        });

        void this.requestAutoDownloadsForRoom(this.watchedRoomId);
      }
    });
  }

  getForMessage(messageId: string): Attachment[] {
    return this.runtimeStore.getAttachmentsForMessage(messageId);
  }

  rememberMessageRoom(messageId: string, roomId: string): void {
    if (!messageId || !roomId)
      return;

    this.runtimeStore.rememberMessageRoom(messageId, roomId);
  }

  queueAutoDownloadsForMessage(messageId: string, attachmentId?: string): void {
    void this.requestAutoDownloadsForMessage(messageId, attachmentId);
  }

  async requestAutoDownloadsForRoom(roomId: string): Promise<void> {
    if (!roomId || !this.isRoomWatched(roomId) || this.webrtc.getConnectedPeers().length === 0)
      return;

    const activeRequest = this.autoDownloadRequestsByRoom.get(roomId);

    if (activeRequest) {
      return activeRequest;
    }

    const request = this.runAutoDownloadsForRoom(roomId).finally(() => {
      if (this.autoDownloadRequestsByRoom.get(roomId) === request) {
        this.autoDownloadRequestsByRoom.delete(roomId);
      }
    });

    this.autoDownloadRequestsByRoom.set(roomId, request);
    return request;
  }

  async deleteForMessage(messageId: string): Promise<void> {
    await this.persistence.deleteForMessage(messageId);
  }

  getAttachmentMetasForMessages(messageIds: string[]): Record<string, AttachmentMeta[]> {
    return this.transfer.getAttachmentMetasForMessages(messageIds);
  }

  async registerSyncedAttachments(
    attachmentMap: Record<string, AttachmentMeta[]>,
    messageRoomIds?: Record<string, string>
  ): Promise<void> {
    await this.transfer.registerSyncedAttachments(attachmentMap, messageRoomIds);

    for (const [messageId, attachments] of Object.entries(attachmentMap)) {
      for (const attachment of attachments) {
        this.queueAutoDownloadsForMessage(messageId, attachment.id);
      }
    }
  }

  requestFromAnyPeer(messageId: string, attachment: Attachment): Promise<void> {
    return this.transfer.requestFromAnyPeer(messageId, attachment);
  }

  handleFileNotFound(payload: FileNotFoundPayload): void {
    this.transfer.handleFileNotFound(payload);
  }

  requestImageFromAnyPeer(messageId: string, attachment: Attachment): Promise<void> {
    return this.transfer.requestImageFromAnyPeer(messageId, attachment);
  }

  hasPendingRequest(messageId: string, attachmentId: string): boolean {
    return this.transfer.hasPendingRequest(messageId, attachmentId);
  }

  async tryRestoreAttachmentFromLocal(attachment: Attachment): Promise<boolean> {
    const restored = await this.persistence.tryRestoreAttachmentFromLocal(attachment);

    if (restored) {
      this.runtimeStore.touch();
    }

    return restored;
  }

  pinDisplayBlobs(attachments: readonly Pick<Attachment, 'id' | 'messageId'>[]): void {
    for (const attachment of attachments) {
      if (!attachment.messageId || !attachment.id) {
        continue;
      }

      this.pinnedDisplayBlobKeys.add(buildAttachmentDisplayPinKey(attachment.messageId, attachment.id));
    }
  }

  unpinDisplayBlobs(attachments: readonly Pick<Attachment, 'id' | 'messageId'>[]): void {
    for (const attachment of attachments) {
      if (!attachment.messageId || !attachment.id) {
        continue;
      }

      this.pinnedDisplayBlobKeys.delete(buildAttachmentDisplayPinKey(attachment.messageId, attachment.id));
    }
  }

  revokeOffscreenDisplayBlobsForMessage(messageId: string): void {
    if (!messageId) {
      return;
    }

    let hasChanges = false;

    for (const attachment of this.runtimeStore.getAttachmentsForMessage(messageId)) {
      if (!shouldRevokeDisplayBlobForAttachment(messageId, attachment, this.pinnedDisplayBlobKeys)) {
        continue;
      }

      if (this.persistence.revokeAttachmentDisplayBlob(attachment)) {
        hasChanges = true;
      }
    }

    if (hasChanges) {
      this.runtimeStore.touch();
    }
  }

  releaseDisplayBlobsForInactiveRooms(activeRoomId: string | null): void {
    const messageIds = collectMessageIdsForInactiveRoomBlobRelease(
      Array.from(this.runtimeStore.getAttachmentEntries(), ([messageId]) => messageId),
      (messageId) => this.runtimeStore.getMessageRoomId(messageId) ?? null,
      activeRoomId
    );

    for (const messageId of messageIds) {
      this.revokeOffscreenDisplayBlobsForMessage(messageId);
    }
  }

  requestFile(messageId: string, attachment: Attachment): Promise<void> {
    return this.transfer.requestFile(messageId, attachment);
  }

  async publishAttachments(
    messageId: string,
    files: File[],
    uploaderPeerId?: string
  ): Promise<void> {
    await this.transfer.publishAttachments(messageId, files, uploaderPeerId);
  }

  handleFileAnnounce(payload: FileAnnouncePayload): void {
    const isNew = this.transfer.handleFileAnnounce(payload);

    if (isNew && payload.messageId && payload.file?.id) {
      this.queueAutoDownloadsForMessage(payload.messageId, payload.file.id);
    }
  }

  handleFileChunk(payload: FileChunkPayload): void {
    this.transfer.handleFileChunk(payload);
  }

  handleFileChunkAck(payload: FileChunkAckPayload): void {
    this.transfer.handleFileChunkAck(payload);
  }

  async handleFileRequest(payload: FileRequestPayload): Promise<void> {
    await this.transfer.handleFileRequest(payload);
  }

  cancelRequest(messageId: string, attachment: Attachment): void {
    this.transfer.cancelRequest(messageId, attachment);
  }

  handleFileCancel(payload: FileCancelPayload): void {
    this.transfer.handleFileCancel(payload);
  }

  async fulfillRequestWithFile(
    messageId: string,
    fileId: string,
    targetPeerId: string,
    file: File
  ): Promise<void> {
    await this.transfer.fulfillRequestWithFile(messageId, fileId, targetPeerId, file);
  }

  private async restoreLocalAttachmentsForRoom(roomId: string): Promise<void> {
    if (!this.isRoomWatched(roomId)) {
      return;
    }

    await this.persistence.whenReady();

    const messageIds = await this.collectMessageIdsForRoom(roomId);

    let hasChanges = false;

    for (const messageId of messageIds) {
      for (const attachment of this.runtimeStore.getAttachmentsForMessage(messageId)) {
        if (await this.persistence.tryRestoreAttachmentHostOnly(attachment)) {
          hasChanges = true;
          await yieldToAttachmentHydrationLoop();
        }
      }
    }

    if (hasChanges) {
      this.runtimeStore.touch();
    }
  }

  private async collectMessageIdsForRoom(roomId: string): Promise<string[]> {
    if (isDirectMessageAttachmentRoomId(roomId)) {
      const messageIds: string[] = [];

      for (const [messageId] of this.runtimeStore.getAttachmentEntries()) {
        const attachmentRoomId = await this.persistence.resolveMessageRoomId(messageId);

        if (attachmentRoomId === roomId) {
          messageIds.push(messageId);
        }
      }

      return messageIds;
    }

    if (!this.database.isReady()) {
      return Array.from(this.runtimeStore.getAttachmentEntries())
        .filter(([messageId]) => this.runtimeStore.getMessageRoomId(messageId) === roomId)
        .map(([messageId]) => messageId);
    }

    const messages = await this.database.getMessages(roomId, 500, 0);

    for (const message of messages) {
      this.runtimeStore.rememberMessageRoom(message.id, message.roomId);
    }

    return messages.map((message) => message.id);
  }

  private async runAutoDownloadsForRoom(roomId: string): Promise<void> {
    if (!this.isRoomWatched(roomId)) {
      return;
    }

    await this.restoreLocalAttachmentsForRoom(roomId);

    let messageIds: string[];

    if (isDirectMessageAttachmentRoomId(roomId)) {
      messageIds = await this.collectMessageIdsForAttachmentsInRoom(roomId);
    } else if (this.database.isReady()) {
      const messages = await this.database.getMessages(roomId, 500, 0);

      for (const message of messages) {
        this.runtimeStore.rememberMessageRoom(message.id, message.roomId);
      }

      messageIds = messages.map((message) => message.id);
    } else {
      messageIds = await this.collectMessageIdsForAttachmentsInRoom(roomId);
    }

    await runTasksWithBoundedConcurrency(
      messageIds.map((messageId) => () => this.requestAutoDownloadsForMessage(messageId)),
      ATTACHMENT_AUTO_DOWNLOAD_CONCURRENCY
    );
  }

  private async collectMessageIdsForAttachmentsInRoom(roomId: string): Promise<string[]> {
    const messageIds: string[] = [];

    for (const [messageId] of this.runtimeStore.getAttachmentEntries()) {
      const attachmentRoomId = await this.persistence.resolveMessageRoomId(messageId);

      if (attachmentRoomId === roomId) {
        messageIds.push(messageId);
      }
    }

    return messageIds;
  }

  private async requestAutoDownloadsForMessage(messageId: string, attachmentId?: string): Promise<void> {
    if (!messageId)
      return;

    const roomId = await this.persistence.resolveMessageRoomId(messageId);

    if (!roomId || !this.isRoomWatched(roomId) || this.webrtc.getConnectedPeers().length === 0) {
      return;
    }

    const attachments = this.runtimeStore.getAttachmentsForMessage(messageId);

    for (const attachment of attachments) {
      if (attachmentId && attachment.id !== attachmentId)
        continue;

      if (!shouldAutoRequestWhenWatched(attachment))
        continue;

      if (attachment.available)
        continue;

      if (shouldResetStalledAttachmentDownload(
        attachment,
        this.transfer.hasPendingRequest(messageId, attachment.id),
        Date.now()
      )) {
        this.transfer.cancelRequest(messageId, attachment);
      } else if ((attachment.receivedBytes ?? 0) > 0) {
        continue;
      }

      if (this.transfer.hasPendingRequest(messageId, attachment.id))
        continue;

      void this.transfer.requestFromAnyPeer(messageId, attachment);
    }
  }

  private extractWatchedRoomId(url: string): string | null {
    return getWatchedAttachmentRoomIdFromUrl(url);
  }

  private async announceHostedAttachments(): Promise<void> {
    const currentUserId = await new Promise<string | null>((resolve) => {
      this.store.select(selectCurrentUserId).pipe(take(1))
        .subscribe((userId) => resolve(userId));
    });

    await this.transfer.reannounceHostedAttachments(currentUserId);
  }

  private isRoomWatched(roomId: string | null | undefined): boolean {
    return !!roomId && roomId === this.watchedRoomId;
  }
}