draw-things-index

/**
 * Document Loader for Draw Things
 * Supports: JSONL (generation logs), PNG/JPG (with embedded metadata)
 */

import { readFile } from 'fs/promises';
import { createHash } from 'crypto';
import path from 'path';
import type { DocumentFormat, ParsedDocument } from '../types';
import { JsonlParser } from './parsers/jsonlParser';
import { PngMetadataParser } from './parsers/pngMetadataParser';

export class DocumentLoader {
  /**
   * Load and parse a document based on its format
   */
  static async load(filePath: string): Promise<ParsedDocument & { hash: string }> {
    const format = this.detectFormat(filePath);
    const hash = await this.hashFile(filePath);

    let parsed: ParsedDocument;

    switch (format) {
      case 'jsonl':
        parsed = await JsonlParser.parse(filePath);
        break;
      
      case 'png':
      case 'jpg':
      case 'jpeg':
      case 'webp':
        parsed = await PngMetadataParser.parse(filePath);
        break;
      
      default:
        throw new Error(`Unsupported file format: ${format}`);
    }

    // Add filename as title if not present
    if (!parsed.metadata.title) {
      parsed.metadata.title = path.basename(filePath, path.extname(filePath));
    }

    return { ...parsed, hash };
  }

  /**
   * Detect document format from file extension
   */
  static detectFormat(filePath: string): DocumentFormat {
    const ext = path.extname(filePath).toLowerCase();
    
    switch (ext) {
      case '.jsonl':
        return 'jsonl';
      case '.png':
        return 'png';
      case '.jpg':
        return 'jpg';
      case '.jpeg':
        return 'jpeg';
      case '.webp':
        return 'webp';
      default:
        return 'unknown';
    }
  }

  /**
   * Calculate SHA-256 hash of file content
   */
  private static async hashFile(filePath: string): Promise<string> {
    const buffer = await readFile(filePath);
    return createHash('sha256').update(buffer).digest('hex');
  }

  /**
   * Check if file format is supported
   */
  static isSupported(filePath: string): boolean {
    const format = this.detectFormat(filePath);
    return ['jsonl', 'png', 'jpg', 'jpeg', 'webp'].includes(format);
  }

  /**
   * Check if file is an image format
   */
  static isImage(filePath: string): boolean {
    const format = this.detectFormat(filePath);
    return ['png', 'jpg', 'jpeg', 'webp'].includes(format);
  }

  /**
   * Check if file is a generation log
   */
  static isGenerationLog(filePath: string): boolean {
    if (this.detectFormat(filePath) !== 'jsonl') return false;
    return path.basename(filePath).toLowerCase() === 'generate-image-plugin.audit.jsonl';
  }
}