import { vectorDatabaseService } from './vectorDatabaseService';
import { llmService } from './llmService';
import { logger } from '../utils/logger';
import { DocumentChunk } from '../models/VectorDatabaseModel';

export interface ChunkingOptions {
  chunkSize: number;
  chunkOverlap: number;
  maxChunks: number;
}

export interface VectorProcessingResult {
  totalChunks: number;
  chunksWithEmbeddings: number;
  processingTime: number;
  averageChunkSize: number;
}

export interface TextBlock {
    type: 'paragraph' | 'table' | 'heading' | 'list_item';
    content: string;
}

export class VectorDocumentProcessor {
  
  /**
   * Store enriched chunks with metadata from agenticRAGProcessor
   */
  async storeDocumentChunks(enrichedChunks: Array<{
    content: string;
    chunkIndex: number;
    startPosition: number;
    endPosition: number;
    sectionType?: string;
    metadata?: {
      hasFinancialData: boolean;
      hasMetrics: boolean;
      keyTerms: string[];
      importance: 'high' | 'medium' | 'low';
      conceptDensity: number;
    };
  }>, options?: {
    documentId: string;
    indexingStrategy?: string;
    similarity_threshold?: number;
    enable_hybrid_search?: boolean;
  }): Promise<void> {
    const startTime = Date.now();
    
    try {
      const documentChunks: DocumentChunk[] = [];
      
      for (const chunk of enrichedChunks) {
        // Generate embedding for the chunk
        const embedding = await vectorDatabaseService.generateEmbeddings(chunk.content);
        
        // Create DocumentChunk with enhanced metadata
        const documentChunk: DocumentChunk = {
          id: `${options?.documentId}-chunk-${chunk.chunkIndex}`,
          documentId: options?.documentId || '',
          content: chunk.content,
          embedding,
          chunkIndex: chunk.chunkIndex,
          metadata: {
            ...chunk.metadata,
            sectionType: chunk.sectionType,
            chunkSize: chunk.content.length,
            processingStrategy: options?.indexingStrategy || 'hierarchical',
            startPosition: chunk.startPosition,
            endPosition: chunk.endPosition
          },
          createdAt: new Date(),
          updatedAt: new Date()
        };
        
        documentChunks.push(documentChunk);
      }
      
      // Store all chunks in vector database
      await vectorDatabaseService.storeDocumentChunks(documentChunks);
      
      const processingTime = Date.now() - startTime;
      const averageImportance = this.calculateAverageImportance(enrichedChunks);
      
      logger.info(`Stored ${documentChunks.length} enriched chunks`, {
        documentId: options?.documentId,
        processingTime,
        averageImportance,
        indexingStrategy: options?.indexingStrategy
      });
      
    } catch (error) {
      logger.error('Failed to store enriched chunks', error);
      throw error;
    }
  }

  /**
   * Calculate average importance score for logging
   */
  private calculateAverageImportance(chunks: Array<{ metadata?: { importance: string } }>): string {
    const importanceScores = chunks
      .map(c => c.metadata?.importance)
      .filter(Boolean);
    
    if (importanceScores.length === 0) return 'unknown';
    
    const highCount = importanceScores.filter(i => i === 'high').length;
    const mediumCount = importanceScores.filter(i => i === 'medium').length;
    
    if (highCount > importanceScores.length / 2) return 'high';
    if (mediumCount + highCount > importanceScores.length / 2) return 'medium';
    return 'low';
  }

  /**
   * Identifies structured blocks of text from a raw string using heuristics.
   * This is the core of the improved ingestion pipeline.
   * @param text The raw text from a PDF extraction.
   */
  private identifyTextBlocks(text: string): TextBlock[] {
    const blocks: TextBlock[] = [];
    // Normalize line endings and remove excessive blank lines to regularize input
    const lines = text.replace(/\n/g, '\n').split('\n');

    let currentParagraph = '';

    for (let i = 0; i < lines.length; i++) {
        const line = lines[i];
        if (line === undefined) continue;
        const trimmedLine = line.trim();

        // If we encounter a blank line, the current paragraph (if any) has ended.
        if (trimmedLine === '') {
            if (currentParagraph.trim()) {
                blocks.push({ type: 'paragraph', content: currentParagraph.trim() });
                currentParagraph = '';
            }
            continue;
        }

        // Heuristic for tables: A line with at least 2 instances of multiple spaces is likely a table row.
        // This is a strong indicator of columnar data in plain text.
        const isTableLike = /(\s{2,}.*){2,}/.test(line);

        if (isTableLike) {
            if (currentParagraph.trim()) {
                blocks.push({ type: 'paragraph', content: currentParagraph.trim() });
                currentParagraph = '';
            }
            // Greedily consume subsequent lines that also look like part of the table.
            let tableContent = line;
            while (i + 1 < lines.length && /(\s{2,}.*){2,}/.test(lines[i + 1] || '')) {
                i++;
                tableContent += '\n' + lines[i];
            }
            blocks.push({ type: 'table', content: tableContent });
            continue;
        }

        // Heuristic for headings: A short line (under 80 chars) that doesn't end with a period.
        // Often in Title Case, but we won't strictly enforce that to be more flexible.
        const isHeadingLike = trimmedLine.length < 80 && !trimmedLine.endsWith('.');
        if (i + 1 < lines.length && (lines[i+1] || '').trim() === '' && isHeadingLike) {
             if (currentParagraph.trim()) {
                blocks.push({ type: 'paragraph', content: currentParagraph.trim() });
                currentParagraph = '';
            }
            blocks.push({ type: 'heading', content: trimmedLine });
            i++; // Skip the blank line after the heading
            continue;
        }

        // Heuristic for list items
        if (trimmedLine.match(/^(\*|-\d+\.)\s/)) {
            if (currentParagraph.trim()) {
                blocks.push({ type: 'paragraph', content: currentParagraph.trim() });
                currentParagraph = '';
            }
            blocks.push({ type: 'list_item', content: trimmedLine });
            continue;
        }

        // Otherwise, append the line to the current paragraph.
        currentParagraph += (currentParagraph ? ' ' : '') + trimmedLine;
    }

    // Add the last remaining paragraph if it exists.
    if (currentParagraph.trim()) {
        blocks.push({ type: 'paragraph', content: currentParagraph.trim() });
    }

    logger.info(`Identified ${blocks.length} semantic blocks from text.`);
    return blocks;
  }

  /**
   * Generates a text summary for a table to be used for embedding.
   * @param tableText The raw text of the table.
   */
  private async getSummaryForTable(tableText: string): Promise<string> {
      const prompt = `The following text is an OCR'd table from a financial document. It may be messy.\n      Summarize the key information in this table in a few clear, narrative sentences.\n      Focus on the main metrics, trends, and time periods.\n      Do not return a markdown table. Return only a natural language summary.\n\n      Table Text:\n      ---\n      ${tableText}\n      ---\n      Summary:`;

      try {
          const result = await llmService.processCIMDocument(prompt, '', { agentName: 'table_summarizer' });
          // Handle both string and object responses from the LLM
          if (result.success) {
              if (typeof result.jsonOutput === 'string') {
                  return result.jsonOutput;
              } 
              if (typeof result.jsonOutput === 'object' && (result.jsonOutput as any)?.summary) {
                  return (result.jsonOutput as any).summary;
              }
          }
          logger.warn('Table summarization failed or returned invalid format, falling back to raw text.', { tableText });
          return tableText; // Fallback
      } catch (error) {
          logger.error('Error during table summarization', { error });
          return tableText; // Fallback
      }
  }

  /**
   * Process document text into chunks and generate embeddings using the new heuristic-based strategy.
   */
  async processDocumentForVectorSearch(
    documentId: string,
    text: string,
    metadata: Record<string, any> = {}
  ): Promise<VectorProcessingResult> {
    const startTime = Date.now();

    try {
      logger.info(`Starting HEURISTIC vector processing for document: ${documentId}`);

      // Step 1: Identify semantic blocks from the document text
      const blocks = this.identifyTextBlocks(text);

      // Step 2: Generate embeddings for each block, with differential processing
      const chunksWithEmbeddings = await this.generateEmbeddingsForBlocks(
        documentId,
        blocks,
        metadata
      );

      // Step 3: Store chunks in vector database
      await vectorDatabaseService.storeDocumentChunks(chunksWithEmbeddings);

      const processingTime = Date.now() - startTime;
      const averageChunkSize = chunksWithEmbeddings.length > 0 ? chunksWithEmbeddings.reduce((sum: number, chunk: any) => sum + chunk.content.length, 0) / chunksWithEmbeddings.length : 0;

      logger.info(`Heuristic vector processing completed for document: ${documentId}`, {
        totalChunks: blocks.length,
        chunksWithEmbeddings: chunksWithEmbeddings.length,
        processingTime,
        averageChunkSize: Math.round(averageChunkSize)
      });

      return {
        totalChunks: blocks.length,
        chunksWithEmbeddings: chunksWithEmbeddings.length,
        processingTime,
        averageChunkSize: Math.round(averageChunkSize)
      };
    } catch (error) {
      logger.error(`Heuristic vector processing failed for document: ${documentId}`, error);
      throw error;
    }
  }

  /**
   * Generates embeddings for the identified text blocks, applying special logic for tables.
   */
  private async generateEmbeddingsForBlocks(
    documentId: string,
    blocks: TextBlock[],
    metadata: Record<string, any>
  ): Promise<DocumentChunk[]> {
    const chunksWithEmbeddings: DocumentChunk[] = [];

    for (let i = 0; i < blocks.length; i++) {
      const block = blocks[i];
      if (!block || !block.content) continue;

      let contentToEmbed = block.content;
      const blockMetadata: any = {
          ...metadata,
          block_type: block.type,
          chunkIndex: i,
          totalChunks: blocks.length,
          chunkSize: block.content.length,
      };

      try {
        // Differential processing for tables
        if (block.type === 'table') {
          logger.info(`Summarizing table chunk ${i}...`);
          contentToEmbed = await this.getSummaryForTable(block.content);
          // Store the original table text in the metadata for later retrieval
          blockMetadata.original_table = block.content;
        }

        const embedding = await vectorDatabaseService.generateEmbeddings(contentToEmbed);

        const documentChunk: DocumentChunk = {
          id: `${documentId}-chunk-${i}`,
          documentId,
          content: contentToEmbed, // This is the summary for tables, or the raw text for others
          metadata: blockMetadata,
          embedding,
          chunkIndex: i,
          createdAt: new Date(),
          updatedAt: new Date()
        };

        chunksWithEmbeddings.push(documentChunk);

        if (blocks.length > 10 && (i + 1) % 10 === 0) {
          logger.info(`Generated embeddings for ${i + 1}/${blocks.length} blocks`);
        }
      } catch (error) {
        logger.error(`Failed to generate embedding for block ${i}`, { error, blockType: block.type });
        // Continue with other chunks, do not halt the entire process
      }
    }

    return chunksWithEmbeddings;
  }

  /**
   * Enhanced search with intelligent filtering and ranking
   */
  async searchRelevantContent(
    query: string,
    options: {
      documentId?: string;
      limit?: number;
      similarityThreshold?: number;
      filters?: Record<string, any>;
      prioritizeFinancial?: boolean;
      boostImportance?: boolean;
      enableReranking?: boolean;
    } = {}
  ) {
    try {
      // Enhanced search parameters
      const searchOptions = {
        ...options,
        limit: Math.min(options.limit || 5, 20), // Cap at 20 for performance
        similarityThreshold: options.similarityThreshold || 0.7, // Higher threshold for quality
      };

      // Add metadata filters for better relevance
      if (options.prioritizeFinancial) {
        searchOptions.filters = {
          ...searchOptions.filters,
          'metadata.hasFinancialData': true
        };
      }

      const rawResults = await vectorDatabaseService.search(query, searchOptions);
      
      // Post-process results for enhanced ranking
      const enhancedResults = this.rankSearchResults(rawResults, query, options);
      
      // Apply reranking if enabled
      let finalResults = enhancedResults;
      if (options.enableReranking !== false) {
        finalResults = await this.rerankResults(query, enhancedResults, options.limit || 5);
      }
      
      logger.info(`Enhanced vector search completed`, {
        query: query.substring(0, 100) + (query.length > 100 ? '...' : ''),
        rawResultsCount: rawResults.length,
        enhancedResultsCount: enhancedResults.length,
        finalResultsCount: finalResults.length,
        documentId: options.documentId,
        prioritizeFinancial: options.prioritizeFinancial,
        enableReranking: options.enableReranking !== false,
        avgRelevanceScore: finalResults.length > 0 ? 
          Math.round((finalResults.reduce((sum: number, r: any) => sum + (r.similarity || 0), 0) / finalResults.length) * 100) / 100 : 0
      });

      return finalResults;
    } catch (error) {
      logger.error('Enhanced vector search failed', { query, options, error });
      throw error;
    }
  }

  /**
   * Rank search results based on multiple criteria
   */
  private rankSearchResults(results: any[], query: string, options: any): any[] {
    return results
      .map(result => ({
        ...result,
        enhancedScore: this.calculateEnhancedScore(result, query, options)
      }))
      .sort((a, b) => b.enhancedScore - a.enhancedScore)
      .slice(0, options.limit || 5);
  }

  /**
   * Calculate enhanced relevance score
   */
  private calculateEnhancedScore(result: any, query: string, options: any): number {
    let score = result.similarity || 0;
    
    // Boost based on importance
    if (options.boostImportance && result.metadata?.importance) {
      if (result.metadata.importance === 'high') score += 0.2;
      else if (result.metadata.importance === 'medium') score += 0.1;
    }
    
    // Boost based on concept density
    if (result.metadata?.conceptDensity) {
      score += result.metadata.conceptDensity * 0.1;
    }
    
    // Boost financial content if query suggests financial context
    if (/financial|revenue|profit|ebitda|margin|cost|cash|debt/i.test(query)) {
      if (result.metadata?.hasFinancialData) score += 0.15;
      if (result.metadata?.hasMetrics) score += 0.1;
    }
    
    // Boost based on section type relevance
    if (result.metadata?.sectionType) {
      const sectionBoosts: Record<string, number> = {
        'executive_summary': 0.1,
        'financial': 0.15,
        'market_analysis': 0.1,
        'management': 0.05
      };
      score += sectionBoosts[result.metadata.sectionType] || 0;
    }
    
    // Boost if query terms appear in key terms
    if (result.metadata?.keyTerms) {
      const queryWords = query.toLowerCase().split(/\s+/);
      const keyTermMatches = result.metadata.keyTerms.filter((term: string) =>
        queryWords.some(word => term.toLowerCase().includes(word))
      ).length;
      score += keyTermMatches * 0.05;
    }
    
    return Math.min(score, 1.0); // Cap at 1.0
  }

  /**
   * Rerank results using cross-encoder approach
   */
  private async rerankResults(query: string, candidates: any[], topK: number = 5): Promise<any[]> {
    try {
      // Create reranking prompt
      const rerankingPrompt = `Given the query: "${query}"

Please rank the following document chunks by relevance (1 = most relevant, ${candidates.length} = least relevant). Consider:
- Semantic similarity to the query
- Financial/business relevance
- Information completeness
- Factual accuracy

Document chunks:
${candidates.map((c, i) => `${i + 1}. ${c.content.substring(0, 200)}...`).join('\n')}

Return only a JSON array of indices in order of relevance: [1, 3, 2, ...]`;

      const result = await llmService.processCIMDocument(rerankingPrompt, '', { 
        agentName: 'reranker',
        maxTokens: 1000
      });

      if (result.success && typeof result.jsonOutput === 'object') {
        const ranking = Array.isArray(result.jsonOutput) ? result.jsonOutput as number[] : null;
        if (ranking) {
          // Apply the ranking
          const reranked = ranking
            .map(index => candidates[index - 1]) // Convert 1-based to 0-based
            .filter(Boolean) // Remove any undefined entries
            .slice(0, topK);
          
          logger.info(`Reranked ${candidates.length} candidates to ${reranked.length} results`);
          return reranked;
        }
      }
      
      // Fallback to original ranking if reranking fails
      logger.warn('Reranking failed, using original ranking');
      return candidates.slice(0, topK);
    } catch (error) {
      logger.error('Reranking failed', error);
      return candidates.slice(0, topK);
    }
  }

  // ... other methods like findSimilarDocuments, etc. remain unchanged ...
}

export const vectorDocumentProcessor = new VectorDocumentProcessor(); 