chunk-storage.service.ts

/**
 * Chunk Storage Service
 * Handles persistence of text chunks and metadata to database
 */

import { TextChunk } from '../types';
import { logger } from '../utils/logger';
import { db } from '../database/connection';

export class ChunkStorageService {
  /**
   * Save chunks to database
   */
  async saveChunks(chunks: TextChunk[]): Promise<void> {
    if (chunks.length === 0) {
      logger.warn('No chunks to save');
      return;
    }

    const documentId = chunks[0].documentId;
    logger.info('Saving chunks to database', {
      documentId,
      chunkCount: chunks.length,
    });

    await db.transaction(async (client) => {
      for (const chunk of chunks) {
        await client.query(
          `INSERT INTO text_chunks 
           (id, document_id, page_number, chunk_index, text, normalized_text, 
            language, token_count, bounding_box, metadata)
           VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
           ON CONFLICT (document_id, chunk_index) 
           DO UPDATE SET
             text = EXCLUDED.text,
             normalized_text = EXCLUDED.normalized_text,
             language = EXCLUDED.language,
             token_count = EXCLUDED.token_count,
             bounding_box = EXCLUDED.bounding_box,
             metadata = EXCLUDED.metadata`,
          [
            chunk.id,
            chunk.documentId,
            chunk.pageNumber,
            chunk.chunkIndex,
            chunk.text,
            chunk.normalizedText,
            chunk.language,
            chunk.tokenCount,
            chunk.boundingBox ? JSON.stringify(chunk.boundingBox) : null,
            chunk.metadata ? JSON.stringify(chunk.metadata) : null,
          ]
        );
      }
    });

    logger.info('Chunks saved successfully', {
      documentId,
      chunkCount: chunks.length,
    });
  }

  /**
   * Retrieve chunks for a document
   */
  async getChunksByDocument(documentId: string): Promise<TextChunk[]> {
    logger.debug('Retrieving chunks for document', { documentId });

    const result = await db.query(
      `SELECT id, document_id, page_number, chunk_index, text, normalized_text,
              language, token_count, bounding_box, metadata, created_at
       FROM text_chunks
       WHERE document_id = $1
       ORDER BY chunk_index ASC`,
      [documentId]
    );

    const chunks: TextChunk[] = result.rows.map((row) => ({
      id: row.id,
      documentId: row.document_id,
      pageNumber: row.page_number,
      chunkIndex: row.chunk_index,
      text: row.text,
      normalizedText: row.normalized_text,
      language: row.language,
      tokenCount: row.token_count,
      boundingBox: row.bounding_box,
      metadata: row.metadata,
    }));

    logger.debug('Retrieved chunks', {
      documentId,
      chunkCount: chunks.length,
    });

    return chunks;
  }

  /**
   * Retrieve chunks for a specific page
   */
  async getChunksByPage(documentId: string, pageNumber: number): Promise<TextChunk[]> {
    logger.debug('Retrieving chunks for page', { documentId, pageNumber });

    const result = await db.query(
      `SELECT id, document_id, page_number, chunk_index, text, normalized_text,
              language, token_count, bounding_box, metadata, created_at
       FROM text_chunks
       WHERE document_id = $1 AND page_number = $2
       ORDER BY chunk_index ASC`,
      [documentId, pageNumber]
    );

    return result.rows.map((row) => ({
      id: row.id,
      documentId: row.document_id,
      pageNumber: row.page_number,
      chunkIndex: row.chunk_index,
      text: row.text,
      normalizedText: row.normalized_text,
      language: row.language,
      tokenCount: row.token_count,
      boundingBox: row.bounding_box,
      metadata: row.metadata,
    }));
  }

  /**
   * Get single chunk by ID
   */
  async getChunkById(chunkId: string): Promise<TextChunk | null> {
    logger.debug('Retrieving chunk by ID', { chunkId });

    const result = await db.query(
      `SELECT id, document_id, page_number, chunk_index, text, normalized_text,
              language, token_count, bounding_box, metadata, created_at
       FROM text_chunks
       WHERE id = $1`,
      [chunkId]
    );

    if (result.rows.length === 0) {
      return null;
    }

    const row = result.rows[0];
    return {
      id: row.id,
      documentId: row.document_id,
      pageNumber: row.page_number,
      chunkIndex: row.chunk_index,
      text: row.text,
      normalizedText: row.normalized_text,
      language: row.language,
      tokenCount: row.token_count,
      boundingBox: row.bounding_box,
      metadata: row.metadata,
    };
  }

  /**
   * Delete chunks for a document
   */
  async deleteChunks(documentId: string): Promise<void> {
    logger.info('Deleting chunks for document', { documentId });

    await db.query('DELETE FROM text_chunks WHERE document_id = $1', [documentId]);

    logger.info('Chunks deleted', { documentId });
  }

  /**
   * Get chunk count for a document
   */
  async getChunkCount(documentId: string): Promise<number> {
    const result = await db.query(
      'SELECT COUNT(*) as count FROM text_chunks WHERE document_id = $1',
      [documentId]
    );

    return parseInt(result.rows[0].count, 10);
  }

  /**
   * Get total token count for a document
   */
  async getTotalTokenCount(documentId: string): Promise<number> {
    const result = await db.query(
      'SELECT COALESCE(SUM(token_count), 0) as total FROM text_chunks WHERE document_id = $1',
      [documentId]
    );

    return parseInt(result.rows[0].total, 10);
  }

  /**
   * Get chunk statistics for a document
   */
  async getChunkStats(documentId: string): Promise<{
    totalChunks: number;
    totalTokens: number;
    avgTokensPerChunk: number;
    minTokens: number;
    maxTokens: number;
  }> {
    logger.debug('Getting chunk statistics', { documentId });

    const result = await db.query(
      `SELECT 
         COUNT(*) as total_chunks,
         COALESCE(SUM(token_count), 0) as total_tokens,
         COALESCE(AVG(token_count), 0) as avg_tokens,
         COALESCE(MIN(token_count), 0) as min_tokens,
         COALESCE(MAX(token_count), 0) as max_tokens
       FROM text_chunks
       WHERE document_id = $1`,
      [documentId]
    );

    const row = result.rows[0];
    return {
      totalChunks: parseInt(row.total_chunks, 10),
      totalTokens: parseInt(row.total_tokens, 10),
      avgTokensPerChunk: parseFloat(row.avg_tokens),
      minTokens: parseInt(row.min_tokens, 10),
      maxTokens: parseInt(row.max_tokens, 10),
    };
  }

  /**
   * Check if chunks exist for document
   */
  async hasChunks(documentId: string): Promise<boolean> {
    const result = await db.query(
      'SELECT EXISTS(SELECT 1 FROM text_chunks WHERE document_id = $1) as exists',
      [documentId]
    );

    return result.rows[0].exists;
  }
}