admin.service.ts

/**
 * Admin Service
 * Handles administrative operations for datasets and documents
 */

import { db } from '../database/connection';
import { logger } from '../utils/logger';
import { QueueService } from './queue.service';
import { ChunkStorageService } from './chunk-storage.service';
import { EmbeddingStorageService } from './embedding-storage.service';
import { TableStorageService } from './table-storage.service';
import { VectorDbService } from './vector-db.service';
import { ocrStorageService } from './ocr-storage.service';

export interface DatasetSummary {
  id: string;
  name: string;
  documentCount: number;
  totalSize: number;
  lastProcessedTime?: Date;
  outstandingJobs: number;
  processingStatus: {
    pending: number;
    processing: number;
    completed: number;
    failed: number;
  };
}

export interface DocumentDetail {
  id: string;
  filename: string;
  fileSize: number;
  mimeType: string;
  uploadTime: Date;
  dataset: string;
  status: string;
  currentStage?: string;
  processingStages: {
    stage: string;
    status: string;
    startedAt?: Date;
    completedAt?: Date;
    errorMessage?: string;
  }[];
  statistics: DocumentStatistics;
}

export interface DocumentStatistics {
  chunkCount?: number;
  embeddingCount?: number;
  tableCount?: number;
  // OCR statistics (Phase 5)
  ocrPageCount?: number;
  ocrAverageConfidence?: number;
  hasStructuredOutput?: boolean;
  hasLayoutData?: boolean;
  hasHTMLOutput?: boolean;
  hasMarkdownOutput?: boolean;
}

export interface JobHistoryEntry {
  jobId: string;
  documentId: string;
  filename: string;
  stage: string;
  status: string;
  attemptsMade: number;
  createdAt: Date;
  completedAt?: Date;
  errorMessage?: string;
}

export class AdminService {
  private queueService: QueueService;
  private chunkStorage: ChunkStorageService;
  private embeddingStorage: EmbeddingStorageService;
  private tableStorage: TableStorageService;
  private vectorDb: VectorDbService;

  constructor() {
    this.queueService = new QueueService();
    this.chunkStorage = new ChunkStorageService();
    this.embeddingStorage = new EmbeddingStorageService();
    this.tableStorage = new TableStorageService();
    this.vectorDb = new VectorDbService();
  }

  /**
   * Get all datasets with summary statistics
   */
  async listDatasets(): Promise<DatasetSummary[]> {
    logger.info('Fetching dataset summaries');

    const result = await db.query(
      `SELECT 
         d.id,
         d.name,
         COUNT(DISTINCT doc.id) as document_count,
         COALESCE(SUM(doc.file_size), 0) as total_size,
         MAX(doc.upload_time) as last_processed_time
       FROM datasets d
       LEFT JOIN documents doc ON d.id = doc.dataset_id
       GROUP BY d.id, d.name
       ORDER BY d.name`
    );

    const datasets: DatasetSummary[] = [];

    for (const row of result.rows) {
      // Get processing status counts
      const statusResult = await db.query(
        `SELECT 
           status,
           COUNT(*) as count
         FROM documents
         WHERE dataset_id = $1
         GROUP BY status`,
        [row.id]
      );

      const processingStatus = {
        pending: 0,
        processing: 0,
        completed: 0,
        failed: 0,
      };

      for (const statusRow of statusResult.rows) {
        const status = statusRow.status as keyof typeof processingStatus;
        if (status in processingStatus) {
          processingStatus[status] = parseInt(statusRow.count);
        }
      }

      // Calculate outstanding jobs (pending + processing)
      const outstandingJobs = processingStatus.pending + processingStatus.processing;

      datasets.push({
        id: row.id,
        name: row.name,
        documentCount: parseInt(row.document_count),
        totalSize: parseInt(row.total_size),
        lastProcessedTime: row.last_processed_time,
        outstandingJobs,
        processingStatus,
      });
    }

    logger.info('Fetched dataset summaries', { count: datasets.length });

    return datasets;
  }

  /**
   * Get detailed information about a specific dataset
   */
  async getDatasetDetails(datasetId: string): Promise<DatasetSummary | null> {
    logger.info('Fetching dataset details', { datasetId });

    const datasets = await this.listDatasets();
    const dataset = datasets.find((d) => d.id === datasetId);

    if (!dataset) {
      logger.warn('Dataset not found', { datasetId });
      return null;
    }

    return dataset;
  }

  /**
   * List all documents with their processing status
   */
  async listDocuments(
    datasetId?: string,
    status?: string,
    limit: number = 50,
    offset: number = 0
  ): Promise<{ documents: DocumentDetail[]; totalCount: number }> {
    logger.info('Listing documents', { datasetId, status, limit, offset });

    // Build WHERE clause
    const conditions: string[] = [];
    const params: any[] = [];
    let paramIndex = 1;

    if (datasetId) {
      conditions.push(`doc.dataset_id = $${paramIndex}`);
      params.push(datasetId);
      paramIndex++;
    }

    if (status) {
      conditions.push(`doc.status = $${paramIndex}`);
      params.push(status);
      paramIndex++;
    }

    const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';

    // Get total count
    const countResult = await db.query(
      `SELECT COUNT(*) as total
       FROM documents doc
       ${whereClause}`,
      params
    );

    const totalCount = parseInt(countResult.rows[0].total);

    // Get documents
    const result = await db.query(
      `SELECT 
         doc.id,
         doc.filename,
         doc.file_size,
         doc.mime_type,
         doc.upload_time,
         d.name as dataset,
         doc.status,
         doc.current_stage
       FROM documents doc
       JOIN datasets d ON doc.dataset_id = d.id
       ${whereClause}
       ORDER BY doc.upload_time DESC
       LIMIT $${paramIndex} OFFSET $${paramIndex + 1}`,
      [...params, limit, offset]
    );

    const documents: DocumentDetail[] = [];

    for (const row of result.rows) {
      const documentDetail = await this.getDocumentDetails(row.id);
      if (documentDetail) {
        documents.push(documentDetail);
      }
    }

    logger.info('Listed documents', { count: documents.length, totalCount });

    return { documents, totalCount };
  }

  /**
   * Get detailed information about a specific document
   */
  async getDocumentDetails(documentId: string): Promise<DocumentDetail | null> {
    logger.info('Fetching document details', { documentId });

    // Get document basic info
    const docResult = await db.query(
      `SELECT 
         doc.id,
         doc.filename,
         doc.file_size,
         doc.mime_type,
         doc.upload_time,
         d.name as dataset,
         doc.status,
         doc.current_stage
       FROM documents doc
       JOIN datasets d ON doc.dataset_id = d.id
       WHERE doc.id = $1`,
      [documentId]
    );

    if (docResult.rows.length === 0) {
      logger.warn('Document not found', { documentId });
      return null;
    }

    const doc = docResult.rows[0];

    // Get processing stages
    const stagesResult = await db.query(
      `SELECT 
         stage,
         status,
         created_at,
         completed_at,
         artifacts
       FROM document_stages
       WHERE document_id = $1
       ORDER BY 
         CASE 
           WHEN stage = 'upload' THEN 1
           WHEN stage = 'language_detection' THEN 2
           WHEN stage = 'ocr' THEN 3
           WHEN stage = 'chunking' THEN 4
           WHEN stage = 'embedding' THEN 5
           WHEN stage = 'table_extraction' THEN 6
           WHEN stage = 'indexing' THEN 7
           ELSE 8
         END`,
      [documentId]
    );

    const processingStages = stagesResult.rows.map((row) => ({
      stage: row.stage,
      status: row.status,
      startedAt: row.created_at,
      completedAt: row.completed_at,
      errorMessage: row.artifacts?.error?.message || null,
    }));

    // Get statistics
    const chunkStats = await this.chunkStorage.getChunkStats(documentId);
    const tableStats = await this.tableStorage.getTableStats(documentId);
    const embeddingCount = await this.embeddingStorage.getEmbeddingCountForDocument(documentId);
    const ocrStats = await ocrStorageService.getOCRStats(documentId);

    const statistics = {
      chunkCount: chunkStats.totalChunks,
      embeddingCount,
      tableCount: tableStats.totalTables,
      // OCR statistics (Phase 5)
      ocrPageCount: ocrStats.pageCount,
      ocrAverageConfidence: ocrStats.averageConfidence,
      hasStructuredOutput: ocrStats.hasLayout || ocrStats.hasHTML || ocrStats.hasMarkdown,
      hasLayoutData: ocrStats.hasLayout,
      hasHTMLOutput: ocrStats.hasHTML,
      hasMarkdownOutput: ocrStats.hasMarkdown,
    };

    const documentDetail: DocumentDetail = {
      id: doc.id,
      filename: doc.filename,
      fileSize: parseInt(doc.file_size),
      mimeType: doc.mime_type,
      uploadTime: doc.upload_time,
      dataset: doc.dataset,
      status: doc.status,
      currentStage: doc.current_stage,
      processingStages,
      statistics,
    };

    logger.info('Fetched document details', { documentId });

    return documentDetail;
  }

  /**
   * Trigger manual reprocessing of a document
   */
  async reprocessDocument(documentId: string, fromStage?: string): Promise<void> {
    logger.info('Triggering document reprocessing', { documentId, fromStage });

    // Verify document exists
    const docResult = await db.query(
      'SELECT id, status FROM documents WHERE id = $1',
      [documentId]
    );

    if (docResult.rows.length === 0) {
      throw new Error('Document not found');
    }

    const currentStatus = docResult.rows[0].status;

    // Don't allow reprocessing of documents currently being processed
    if (currentStatus === 'processing') {
      throw new Error('Document is currently being processed. Please wait for completion or failure.');
    }

    // Determine which stage to restart from
    const restartStage = fromStage || 'language_detection'; // Default to language detection

    // Reset document status
    await db.query(
      'UPDATE documents SET status = $1, current_stage = $2 WHERE id = $3',
      ['processing', restartStage, documentId]
    );

    // Reset stages from restart point onwards
    const stageOrder = [
      'upload',
      'language_detection',
      'ocr',
      'chunking',
      'embedding',
      'table_extraction',
      'indexing',
    ];

    const restartIndex = stageOrder.indexOf(restartStage);

    if (restartIndex === -1) {
      throw new Error(`Invalid stage: ${restartStage}`);
    }

    const stagesToReset = stageOrder.slice(restartIndex);

    for (const stage of stagesToReset) {
      await db.query(
        `UPDATE document_stages 
         SET status = $1, completed_at = NULL, artifacts = NULL, updated_at = CURRENT_TIMESTAMP
         WHERE document_id = $2 AND stage = $3`,
        ['pending', documentId, stage]
      );
    }

    // Queue the document for processing starting from restart stage
    await this.queueService.addJob(`processing-${restartStage}`, {
      documentId,
      stage: restartStage,
      timestamp: new Date().toISOString(),
    });

    logger.info('Document reprocessing triggered', { documentId, restartStage });
  }

  /**
   * Delete a document and all its associated data
   */
  async deleteDocument(documentId: string): Promise<void> {
    logger.info('Deleting document', { documentId });

    // Start transaction
    const client = await db.getClient();

    try {
      await client.query('BEGIN');

      // Delete from vector database
      await this.vectorDb.deleteDocumentVectors(documentId);

      // Delete embeddings
      await client.query('DELETE FROM chunk_embeddings WHERE chunk_id IN (SELECT id FROM text_chunks WHERE document_id = $1)', [documentId]);

      // Delete chunks
      await client.query('DELETE FROM text_chunks WHERE document_id = $1', [documentId]);

      // Delete tables
      await client.query('DELETE FROM table_cells WHERE table_id IN (SELECT id FROM extracted_tables WHERE document_id = $1)', [documentId]);
      await client.query('DELETE FROM extracted_tables WHERE document_id = $1', [documentId]);

      // Delete OCR results
      await client.query('DELETE FROM ocr_results WHERE document_id = $1', [documentId]);

      // Delete processing stages
      await client.query('DELETE FROM document_stages WHERE document_id = $1', [documentId]);

      // Delete document
      await client.query('DELETE FROM documents WHERE id = $1', [documentId]);

      await client.query('COMMIT');

      logger.info('Document deleted successfully', { documentId });
    } catch (error) {
      await client.query('ROLLBACK');
      logger.error('Failed to delete document', { error, documentId });
      throw error;
    } finally {
      client.release();
    }
  }

  /**
   * Get recent job history
   */
  async getJobHistory(
    limit: number = 100,
    offset: number = 0
  ): Promise<{ jobs: JobHistoryEntry[]; totalCount: number }> {
    logger.info('Fetching job history', { limit, offset });

    // Get job history from processing_jobs table
    const countResult = await db.query(
      'SELECT COUNT(*) as total FROM processing_jobs'
    );

    const totalCount = parseInt(countResult.rows[0].total);

    const result = await db.query(
      `SELECT 
         pj.id as job_id,
         pj.document_id,
         d.filename,
         pj.stage,
         pj.status,
         pj.retry_count as attempts_made,
         pj.start_time as created_at,
         pj.end_time as completed_at,
         pj.error_details->>'message' as error_message
       FROM processing_jobs pj
       JOIN documents d ON pj.document_id = d.id
       ORDER BY pj.start_time DESC
       LIMIT $1 OFFSET $2`,
      [limit, offset]
    );

    const jobs: JobHistoryEntry[] = result.rows.map((row) => ({
      jobId: row.job_id,
      documentId: row.document_id,
      filename: row.filename,
      stage: row.stage,
      status: row.status,
      attemptsMade: row.attempts_made,
      createdAt: row.created_at,
      completedAt: row.completed_at,
      errorMessage: row.error_message,
    }));

    logger.info('Fetched job history', { count: jobs.length, totalCount });

    return { jobs, totalCount };
  }

  /**
   * Get system statistics
   */
  async getSystemStats(): Promise<{
    totalDocuments: number;
    totalDatasets: number;
    processingStatus: {
      pending: number;
      processing: number;
      completed: number;
      failed: number;
    };
    storageUsed: number;
    vectorCount: number;
  }> {
    logger.info('Fetching system statistics');

    // Total documents
    const docCountResult = await db.query('SELECT COUNT(*) as total FROM documents');
    const totalDocuments = parseInt(docCountResult.rows[0].total);

    // Total datasets
    const datasetCountResult = await db.query('SELECT COUNT(*) as total FROM datasets');
    const totalDatasets = parseInt(datasetCountResult.rows[0].total);

    // Processing status counts
    const statusResult = await db.query(
      `SELECT status, COUNT(*) as count
       FROM documents
       GROUP BY status`
    );

    const processingStatus = {
      pending: 0,
      processing: 0,
      completed: 0,
      failed: 0,
    };

    for (const row of statusResult.rows) {
      const status = row.status as keyof typeof processingStatus;
      if (status in processingStatus) {
        processingStatus[status] = parseInt(row.count);
      }
    }

    // Storage used (sum of file sizes)
    const storageResult = await db.query('SELECT COALESCE(SUM(file_size), 0) as total FROM documents');
    const storageUsed = parseInt(storageResult.rows[0].total);

    // Vector count
    let vectorCount = 0;
    try {
      const collectionInfo = await this.vectorDb.getCollectionInfo();
      vectorCount = collectionInfo.vectorsCount;
    } catch (error) {
      logger.warn('Failed to get vector count', { error });
    }

    const stats = {
      totalDocuments,
      totalDatasets,
      processingStatus,
      storageUsed,
      vectorCount,
    };

    logger.info('Fetched system statistics', stats);

    return stats;
  }
}