import { VectorDocumentProcessor, TextBlock } from '../vectorDocumentProcessor';
import { llmService } from '../llmService';
import { vectorDatabaseService } from '../vectorDatabaseService';

// Mock the dependencies
jest.mock('../llmService');
jest.mock('../vectorDatabaseService');

const mockedLlmService = llmService as jest.Mocked<typeof llmService>;
const mockedVectorDBService = vectorDatabaseService as jest.Mocked<typeof vectorDatabaseService>;

// Sample text mimicking a messy PDF extraction with various elements
const sampleText = `
This is the first paragraph of the document. It contains some general information about the company.

Financial Highlights

This paragraph discusses the financial performance. It is located after a heading.

Here is a table of financial data:

Metric      FY2021      FY2022      FY2023
Revenue     $10.0M      $12.5M      $15.0M
EBITDA      $2.0M       $2.5M       $3.0M

This is the final paragraph, coming after the table. It summarizes the outlook.
`;

describe('VectorDocumentProcessor', () => {
  let processor: VectorDocumentProcessor;

  beforeEach(() => {
    processor = new VectorDocumentProcessor();
    // Reset mocks before each test
    jest.clearAllMocks();
    
    // Set up VectorDatabaseService mock methods
    (mockedVectorDBService as any).generateEmbeddings = jest.fn();
    (mockedVectorDBService as any).storeDocumentChunks = jest.fn();
    (mockedVectorDBService as any).search = jest.fn();
  });

  describe('identifyTextBlocks', () => {
    it('should correctly identify paragraphs, headings, and tables', () => {
      // Access the private method for testing purposes
      const blocks: TextBlock[] = (processor as any).identifyTextBlocks(sampleText);

      expect(blocks).toHaveLength(5);
      
      // Check block types
      expect(blocks[0]?.type).toBe('paragraph');
      expect(blocks[1]?.type).toBe('heading');
      expect(blocks[2]?.type).toBe('paragraph');
      expect(blocks[3]?.type).toBe('table');
      expect(blocks[4]?.type).toBe('paragraph');

      // Check block content
      expect(blocks[0]?.content).toBe('This is the first paragraph of the document. It contains some general information about the company.');
      expect(blocks[1]?.content).toBe('Financial Highlights');
      expect(blocks[3]?.content).toContain('Revenue     $10.0M      $12.5M      $15.0M');
    });
  });

  describe('processDocumentForVectorSearch', () => {
    it('should use the LLM to summarize tables and store original in metadata', async () => {
      const documentId = 'test-doc-1';
      const tableSummary = 'The table shows revenue growing from $10M to $15M and EBITDA growing from $2M to $3M between FY2021 and FY2023.';

      // Mock the LLM service to return a summary for the table
      mockedLlmService.processCIMDocument.mockResolvedValue({
        success: true,
        jsonOutput: { summary: tableSummary } as any,
        model: 'test-model',
        cost: 0.01,
        inputTokens: 100,
        outputTokens: 50,
      });

      // Mock the embedding service to return a dummy vector
      (mockedVectorDBService as any).generateEmbeddings.mockResolvedValue([0.1, 0.2, 0.3]);

      // Mock the storage service
      (mockedVectorDBService as any).storeDocumentChunks.mockResolvedValue();

      await processor.processDocumentForVectorSearch(documentId, sampleText);

      // Verify that storeDocumentChunks was called
      expect((mockedVectorDBService as any).storeDocumentChunks).toHaveBeenCalled();

      // Get the arguments passed to storeDocumentChunks
      const storedChunks = (mockedVectorDBService as any).storeDocumentChunks.mock.calls[0]?.[0];
      expect(storedChunks).toBeDefined();
      if (!storedChunks) return;

      expect(storedChunks).toHaveLength(5);

      // Find the table chunk
      const tableChunk = storedChunks.find((c: any) => c.metadata.block_type === 'table');
      expect(tableChunk).toBeDefined();
      if (!tableChunk) return;

      // Assert that the LLM was called for the table summarization
      expect(mockedLlmService.processCIMDocument).toHaveBeenCalledTimes(1);
      const prompt = mockedLlmService.processCIMDocument.mock.calls[0]?.[0];
      expect(prompt).toContain('Summarize the key information in this table');

      // Assert that the table chunk's content is the LLM summary
      expect(tableChunk.content).toBe(tableSummary);

      // Assert that the original table text is stored in the metadata
      expect(tableChunk.metadata['original_table']).toContain('Metric      FY2021      FY2022      FY2023');

      // Find a paragraph chunk and check its content
      const paragraphChunk = storedChunks.find((c: any) => c.metadata.block_type === 'paragraph');
      expect(paragraphChunk).toBeDefined();
      if (paragraphChunk) {
        expect(paragraphChunk.content).not.toBe(tableSummary); // Ensure it wasn't summarized
      }
    });
  });
});
