RedditBot/src/services/media.ts

import fs from 'fs/promises';
import path from 'path';
import crypto from 'crypto';
import { z } from 'zod';
import sharp from 'sharp';
import type { Logger } from '../utils/logger.js';
import type { ActiveChat, ChatMedia, ChatMessage } from './chat.js';
import type { AccountContentRepository } from '../database/account-repositories.js';

export interface MediaDownloadResult {
  buffer: Buffer;
  contentType: string | null;
}

export interface MediaDownloader {
  fetch(url: string): Promise<MediaDownloadResult | null>;
}

export interface MediaRecord {
  accountId?: string;
  id: string;
  chatUsername: string | null;
  chatRoomId: string | null;
  sender: string;
  timestamp: number;
  originalUrl: string;
  fileName: string;
  relativePath: string;
  createdAt: number;
}

export class MediaAttachmentError extends Error {
  constructor(public readonly failedDownloads: number) {
    super(`Failed to retrieve ${failedDownloads} conversation media item(s)`);
    this.name = 'MediaAttachmentError';
  }
}

class MediaDownloadError extends Error {}

const MediaIndexSchema = z.object({
  version: z.number(),
  savedAt: z.number(),
  items: z.array(
    z.object({
      id: z.string(),
      chatUsername: z.string().nullable(),
      chatRoomId: z.string().nullable(),
      sender: z.string(),
      timestamp: z.number(),
      originalUrl: z.string(),
      fileName: z.string(),
      relativePath: z.string(),
      createdAt: z.number(),
    })
  ),
});

type MediaIndex = z.infer<typeof MediaIndexSchema>;

interface ScrapedChatMessage extends ChatMessage {
  mediaUrls: string[];
}

export class MediaService {
  private readonly mediaRoot: string;
  private readonly indexPath: string;
  private index: MediaIndex | null = null;

  constructor(
    private downloader: MediaDownloader,
    private logger: Logger,
    mediaRoot?: string,
    private readonly databaseStore?: AccountContentRepository,
    private readonly getAccountId: () => string | null = () => null,
  ) {
    this.mediaRoot = mediaRoot ?? path.join(process.cwd(), 'data', 'media');
    this.indexPath = path.join(this.mediaRoot, 'index.json');
  }

  async attachMedia(chat: ActiveChat, messages: ScrapedChatMessage[]): Promise<ChatMessage[]> {
    const databaseBacked = Boolean(this.databaseStore && this.getAccountId());
    if (!databaseBacked) await this.loadIndex();
    const updated: ChatMessage[] = [];
    let failedDownloads = 0;

    for (const message of messages) {
      let media: ChatMedia[] = [];
      try {
        media = await this.downloadMessageMedia(chat, message);
      } catch (error) {
        if (!(error instanceof MediaDownloadError)) throw error;
        failedDownloads++;
      }
      updated.push({
        eventId: message.eventId,
        sender: message.sender,
        text: message.text,
        timestamp: message.timestamp,
        isSelf: message.isSelf,
        media: media.length > 0 ? media : undefined,
      });
    }

    if (!databaseBacked) await this.saveIndex();
    if (failedDownloads > 0) throw new MediaAttachmentError(failedDownloads);
    return updated;
  }

  async getMediaItems(): Promise<MediaRecord[]> {
    const accountId = this.getAccountId();
    if (this.databaseStore && accountId) {
      return this.databaseStore.listMedia(accountId).map((item) => ({
        accountId: item.accountId,
        id: item.mediaId,
        chatUsername: item.chatUsername,
        chatRoomId: item.chatRoomId,
        sender: item.sender,
        timestamp: item.timestamp,
        originalUrl: item.originalUrl,
        fileName: item.fileName,
        relativePath: item.relativePath,
        createdAt: item.createdAt,
      }));
    }
    await this.loadIndex();
    return (this.index?.items || []).slice().sort((a, b) => b.timestamp - a.timestamp);
  }

  private async downloadMessageMedia(chat: ActiveChat, message: ScrapedChatMessage): Promise<ChatMedia[]> {
    if (!message.mediaUrls || message.mediaUrls.length === 0) {
      return [];
    }

    const mediaItems: ChatMedia[] = [];

    for (const url of message.mediaUrls) {
      const record = await this.downloadMedia(chat, message, url);
      if (record) {
        mediaItems.push({
          originalUrl: record.originalUrl,
          localPath: record.relativePath,
          fileName: record.fileName,
        });
      }
    }

    return mediaItems;
  }

  /**
   * Convert Reddit Matrix blurred thumbnail URLs to full-resolution download URLs.
   * Blurred: https://matrix.redditspace.com/_matrix/media/v3/thumbnail/reddit.com/blur_XXXXX?height=300
   * Full:    https://matrix.redditspace.com/_matrix/media/v3/download/reddit.com/XXXXX
   */
  private resolveFullMediaUrl(url: string): string {
    try {
      const parsed = new URL(url);
      if (parsed.hostname === 'matrix.redditspace.com' && parsed.pathname.includes('/thumbnail/')) {
        // Strip blur_ prefix from the media ID
        const fullPath = parsed.pathname.replace('/thumbnail/', '/download/').replace('/blur_', '/');
        return `${parsed.origin}${fullPath}`;
      }
    } catch { /* keep original */ }
    return url;
  }

  private async downloadMedia(chat: ActiveChat, message: ScrapedChatMessage, url: string): Promise<MediaRecord | null> {
    if (!url || url.startsWith('data:')) {
      return null;
    }

    // Try the full-resolution URL first, fall back to original
    const fullUrl = this.resolveFullMediaUrl(url);
    const accountId = this.getAccountId();
    const ownerKey = accountId ?? 'legacy';
    const assetId = crypto.createHash('sha1').update(`${ownerKey}\0${fullUrl}`).digest('hex');
    const occurrenceKey = message.eventId
      ?? `${chat.roomId ?? ''}\0${message.timestamp}\0${message.sender}`;
    const id = crypto.createHash('sha1')
      .update(`${ownerKey}\0${occurrenceKey}\0${fullUrl}`)
      .digest('hex');
    if (this.databaseStore && accountId) {
      const existingRecord = this.databaseStore.getMedia(accountId, id);
      if (existingRecord) {
        try {
          await fs.access(path.join(this.mediaRoot, existingRecord.relativePath));
          return { ...existingRecord, id: existingRecord.mediaId };
        } catch {
          this.databaseStore.deleteMedia(accountId, id);
        }
      }

      const reusable = this.databaseStore.getMediaByOriginalUrl(accountId, fullUrl);
      if (reusable) {
        try {
          await fs.access(path.join(this.mediaRoot, reusable.relativePath));
          const record = this.buildRecord(
            id, chat, message, fullUrl, reusable.fileName, reusable.relativePath, accountId,
          );
          this.databaseStore.upsertMedia(this.toStoredRecord(record, accountId));
          if (reusable.mediaId === assetId && assetId !== id) {
            this.databaseStore.deleteMedia(accountId, assetId);
          }
          return record;
        } catch {
          this.databaseStore.deleteMedia(accountId, reusable.mediaId);
        }
      }
    }
    const existing = this.index?.items.find((item) => item.id === id);
    if (existing) {
      // Verify the file still exists on disk
      const existingPath = path.join(this.mediaRoot, existing.relativePath);
      try {
        await fs.access(existingPath);
        return existing;
      } catch {
        // File missing — remove stale record and re-download
        this.logger.debug(`Media file missing on disk, re-downloading: ${existing.fileName}`);
        this.index!.items = this.index!.items.filter((item) => item.id !== id);
      }
    }

    const reusable = this.index?.items.find((item) => item.originalUrl === fullUrl);
    if (reusable) {
      try {
        await fs.access(path.join(this.mediaRoot, reusable.relativePath));
        const record = this.buildRecord(
          id, chat, message, fullUrl, reusable.fileName, reusable.relativePath, accountId,
        );
        this.index!.items.push(record);
        if (reusable.id === assetId && assetId !== id) {
          this.index!.items = this.index!.items.filter((item) => item.id !== assetId);
        }
        return record;
      } catch {
        this.index!.items = this.index!.items.filter((item) => item.id !== reusable.id);
      }
    }

    let download = await this.downloader.fetch(fullUrl);
    let extension = download
      ? await this.pickExtension(url, download.contentType, download.buffer)
      : null;
    if ((!download || !extension) && fullUrl !== url) {
      this.logger.debug(`Full-res media was unavailable or invalid, trying original URL: ${url}`);
      download = await this.downloader.fetch(url);
      extension = download
        ? await this.pickExtension(url, download.contentType, download.buffer)
        : null;
    }
    if (!download) {
      this.logger.warn(`Failed to download chat media from ${fullUrl}`);
      throw new MediaDownloadError('Conversation media download failed');
    }

    if (!extension) {
      this.logger.debug(`Skipping non-image media from ${url}`);
      return null;
    }

    const fileName = `${assetId}${extension}`;
    const accountDirectory = accountId
      ? crypto.createHash('sha256').update(accountId).digest('hex').slice(0, 16)
      : '';
    const relativePath = accountDirectory ? path.join(accountDirectory, fileName) : fileName;
    const fullPath = path.join(this.mediaRoot, relativePath);

    await fs.mkdir(path.dirname(fullPath), { recursive: true });
    await fs.writeFile(fullPath, download.buffer);

    const record = this.buildRecord(id, chat, message, fullUrl, fileName, relativePath, accountId);

    if (this.databaseStore && accountId) {
      this.databaseStore.upsertMedia(this.toStoredRecord(record, accountId));
      return record;
    }

    if (!this.index) {
      this.index = { version: 1, savedAt: Date.now(), items: [record] };
    } else {
      this.index.items.push(record);
      this.index.savedAt = Date.now();
    }

    return record;
  }

  private buildRecord(
    id: string,
    chat: ActiveChat,
    message: ScrapedChatMessage,
    originalUrl: string,
    fileName: string,
    relativePath: string,
    accountId: string | null,
  ): MediaRecord {
    return {
      accountId: accountId ?? undefined,
      id,
      chatUsername: chat.username || null,
      chatRoomId: chat.roomId || null,
      sender: message.sender,
      timestamp: message.timestamp,
      originalUrl,
      fileName,
      relativePath,
      createdAt: Date.now(),
    };
  }

  private toStoredRecord(record: MediaRecord, accountId: string) {
    return {
      accountId,
      mediaId: record.id,
      chatUsername: record.chatUsername,
      chatRoomId: record.chatRoomId,
      sender: record.sender,
      timestamp: record.timestamp,
      originalUrl: record.originalUrl,
      fileName: record.fileName,
      relativePath: record.relativePath,
      createdAt: record.createdAt,
    };
  }

  private async pickExtension(url: string, contentType: string | null, buffer: Buffer): Promise<string | null> {
    const supportedExtensions = new Set(['.jpg', '.jpeg', '.png', '.gif', '.webp', '.avif']);
    let urlExt = '';
    try {
      urlExt = path.extname(new URL(url).pathname || '').toLowerCase();
    } catch {
      urlExt = path.extname(url.split('?')[0] || '').toLowerCase();
    }
    if (supportedExtensions.has(urlExt)) {
      return urlExt;
    }

    if (contentType?.includes('image/jpeg')) return '.jpg';
    if (contentType?.includes('image/png')) return '.png';
    if (contentType?.includes('image/gif')) return '.gif';
    if (contentType?.includes('image/webp')) return '.webp';
    if (contentType?.includes('image/avif')) return '.avif';

    try {
      const format = (await sharp(buffer).metadata()).format;
      if (format === 'jpeg') return '.jpg';
      if (format === 'png') return '.png';
      if (format === 'gif') return '.gif';
      if (format === 'webp') return '.webp';
      if (format === 'heif') return '.avif';
    } catch {
      // Unsupported/corrupt payloads are not gallery images.
    }

    return null;
  }

  private async loadIndex(): Promise<void> {
    if (this.index) {
      return;
    }

    try {
      const raw = await fs.readFile(this.indexPath, 'utf-8');
      const parsed = MediaIndexSchema.safeParse(JSON.parse(raw));
      if (!parsed.success) {
        this.logger.warn('Media index is invalid, starting fresh.');
        this.index = { version: 1, savedAt: Date.now(), items: [] };
        return;
      }

      this.index = parsed.data;
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code !== 'ENOENT') {
        this.logger.warn('Failed to load media index, starting fresh:', error);
      }
      this.index = { version: 1, savedAt: Date.now(), items: [] };
    }
  }

  private async saveIndex(): Promise<void> {
    if (!this.index) {
      return;
    }

    try {
      await fs.mkdir(this.mediaRoot, { recursive: true });
      await fs.writeFile(this.indexPath, JSON.stringify(this.index, null, 2), 'utf-8');
    } catch (error) {
      this.logger.warn('Failed to save media index:', error);
    }
  }
}