draw-things-index

/**
 * JSONL Parser for Draw Things / LM Studio Generation Audit Logs
 * Parses .jsonl files containing image generation metadata
 */

import { readFile } from 'fs/promises';
import path from 'path';
import type { ParsedDocument } from '../../types';

const SUPPORTED_AUDIT_JSONL_BASENAME = 'generate-image-plugin.audit.jsonl';

function assertSupportedAuditJsonlPath(filePath: string): void {
  const base = path.basename(filePath).toLowerCase();
  if (base !== SUPPORTED_AUDIT_JSONL_BASENAME) {
    throw new Error(
      `[JsonlParser] Unsupported .jsonl file: ${filePath}. Only ${SUPPORTED_AUDIT_JSONL_BASENAME} is supported.`,
    );
  }
}

/**
 * Structure of a generation log entry from generate-image-plugin.audit.jsonl
 */
export interface GenerationLogEntry {
  timestamp: string;
  requestId: string;
  backend: string;
  mode: string;
  chat_id: string;
  user_request: {
    prompt: string;
    mode: string;
    model: string;
    width: number;
    height: number;
    imageFormat?: string;
    variants?: number;
  };
  render_target?: {
    requested_raw: { width: number; height: number };
    requested_effective: { width: number; height: number };
    needs_upscaler: boolean;
  };
  output: {
    backend_returned?: { width: number; height: number };
    post_processed?: { width: number; height: number };
    inference_time_ms: number;
    model_used: string;
    overlay_source?: string;
    overlay_preset?: string;
    loras_used?: string[];
    prompt_used: string;
    prompt_origin: string;
    variants: Array<{
      v: number;
      path: string;
      url: string;
      bytes: number;
      preview_path?: string;
      preview_url?: string;
      http_preview_url?: string;
    }>;
  };
}

/**
 * Parsed generation with searchable text
 */
export interface ParsedGeneration {
  entry: GenerationLogEntry;
  searchableText: string;
  imagePaths: string[];
  httpPreviewUrls: string[];
}

export class JsonlParser {
  /**
   * Parse JSONL file containing generation logs
   */
  static async parse(filePath: string): Promise<ParsedDocument> {
    assertSupportedAuditJsonlPath(filePath);
    const content = await readFile(filePath, 'utf-8');
    const generations = this.parseJsonlContent(content);
    
    // Create searchable text from all generations
    const searchableContent = generations
      .map(gen => gen.searchableText)
      .join('\n\n---\n\n');
    
    return {
      content: searchableContent,
      metadata: {
        format: 'jsonl',
        generationCount: generations.length,
        generations: generations.map(g => ({
          timestamp: g.entry.timestamp,
          prompt: g.entry.output.prompt_used,
          model: g.entry.output.model_used,
          loras: g.entry.output.loras_used,
          inferenceTimeMs: g.entry.output.inference_time_ms,
          imagePaths: g.imagePaths,
        })),
      },
    };
  }

  /**
   * Parse JSONL content into generation entries
   * Supports both true JSONL (one JSON per line) and pretty-printed JSON objects
   */
  static parseJsonlContent(content: string): ParsedGeneration[] {
    const generations: ParsedGeneration[] = [];
    
    // First, try to parse as array of objects (pretty-printed)
    const trimmed = content.trim();
    
    // Check if it's a JSON array
    if (trimmed.startsWith('[')) {
      try {
        const entries = JSON.parse(trimmed) as GenerationLogEntry[];
        for (const entry of entries) {
          if (entry.output) {
            generations.push(this.createParsedGeneration(entry));
          }
        }
        return generations;
      } catch {
        // Not a valid array, try other methods
      }
    }
    
    // Try to parse as concatenated JSON objects (pretty-printed)
    // Split on }{ or }\n{ patterns
    const jsonObjects = this.splitJsonObjects(trimmed);
    
    for (const jsonStr of jsonObjects) {
      try {
        const entry = JSON.parse(jsonStr) as GenerationLogEntry;
        if (entry.output) {
          generations.push(this.createParsedGeneration(entry));
        }
      } catch (e) {
        // Skip malformed entries silently unless debugging
      }
    }

    // Fallback: try line-by-line for true JSONL
    if (generations.length === 0) {
      const lines = content.split('\n').filter(line => line.trim());
      for (const line of lines) {
        try {
          const entry = JSON.parse(line) as GenerationLogEntry;
          if (entry.output) {
            generations.push(this.createParsedGeneration(entry));
          }
        } catch {
          // Skip malformed lines
        }
      }
    }

    return generations;
  }

  /**
   * Split concatenated JSON objects from pretty-printed content
   */
  private static splitJsonObjects(content: string): string[] {
    const objects: string[] = [];
    let depth = 0;
    let start = -1;
    
    for (let i = 0; i < content.length; i++) {
      const char = content[i];
      
      if (char === '{') {
        if (depth === 0) {
          start = i;
        }
        depth++;
      } else if (char === '}') {
        depth--;
        if (depth === 0 && start !== -1) {
          objects.push(content.slice(start, i + 1));
          start = -1;
        }
      }
    }
    
    return objects;
  }

  /**
   * Create searchable text from a generation entry
   */
  static createParsedGeneration(entry: GenerationLogEntry): ParsedGeneration {
    const imagePaths = entry.output?.variants?.map(v => v.path) || [];
    const httpPreviewUrls = entry.output?.variants
      ?.map(v => v.http_preview_url)
      .filter((url): url is string => !!url) || [];
    
    // Build searchable text with all relevant metadata
    const parts: string[] = [
      `Prompt: ${entry.output.prompt_used}`,
    ];

    if (entry.user_request.prompt !== entry.output.prompt_used) {
      parts.push(`Original Request: ${entry.user_request.prompt}`);
    }

    parts.push(`Model: ${entry.output.model_used}`);
    
    if (entry.output.loras_used?.length) {
      parts.push(`LoRAs: ${entry.output.loras_used.join(', ')}`);
    }

    parts.push(`Mode: ${entry.mode}`);
    parts.push(`Size: ${entry.render_target?.requested_effective?.width || entry.user_request.width}x${entry.render_target?.requested_effective?.height || entry.user_request.height}`);
    parts.push(`Inference Time: ${(entry.output.inference_time_ms / 1000).toFixed(1)}s`);
    parts.push(`Generated: ${entry.timestamp}`);
    
    if (imagePaths.length) {
      parts.push(`Images: ${imagePaths.length} variant(s)`);
      imagePaths.forEach((p, i) => {
        parts.push(`  [${i + 1}] ${p}`);
      });
    }

    return {
      entry,
      searchableText: parts.join('\n'),
      imagePaths,
      httpPreviewUrls,
    };
  }

  /**
   * Parse and return individual generations (for chunking per generation)
   */
  static async parseAsGenerations(filePath: string): Promise<ParsedGeneration[]> {
    assertSupportedAuditJsonlPath(filePath);
    const content = await readFile(filePath, 'utf-8');
    return this.parseJsonlContent(content);
  }
}