"""
Comprehensive test suite for Word document extraction with V2 pipeline.

Tests all 51 Word files in CustomerDocument/ folder to validate:
- File discovery and tracking
- Format normalization (.doc/.docm → .docx)
- Word → DoclingDocument conversion
- Semantic chunking with HybridChunker
- Success rate ≥85% requirement
- Metrics generation by type and format
"""

import pytest
from pathlib import Path
from uuid import uuid4
from typing import List, Dict, Any

from src.extraction_v2.pipeline import ExtractionPipelineV2
from batch_extract_v2 import BatchProcessor


class TestWordFileDiscovery:
    """Test Word file discovery in batch processor."""

    def test_find_all_files_discovers_word_files(self):
        """Test that batch processor discovers all Word files."""
        processor = BatchProcessor("CustomerDocument/")
        all_files = processor.find_all_files()

        # Filter for Word files
        word_files = [f for f in all_files if f.suffix.lower() in [".docx", ".doc", ".docm"]]

        # Filter out preprocessed files (created during testing)
        word_files = [f for f in word_files if "_preprocessed" not in f.name]

        # Verify we found Word files
        assert len(word_files) > 0, "No Word files found in CustomerDocument/"

        print(f"\nFound {len(word_files)} Word files (excluding preprocessed):")
        print(f"  Total: {len(word_files)}")

        # Count by format
        docx_count = sum(1 for f in word_files if f.suffix.lower() == ".docx")
        doc_count = sum(1 for f in word_files if f.suffix.lower() == ".doc")
        docm_count = sum(1 for f in word_files if f.suffix.lower() == ".docm")

        print(f"  .docx: {docx_count}")
        print(f"  .doc: {doc_count}")
        print(f"  .docm: {docm_count}")

        # Verify expected distribution (from tech-spec: 2 docx, 21 doc, 28 docm = 51 total)
        assert len(word_files) >= 51, f"Expected at least 51 Word files, found {len(word_files)}"
        assert doc_count == 21, f"Expected 21 .doc files, found {doc_count}"
        assert docm_count == 28, f"Expected 28 .docm files, found {docm_count}"

    def test_categorize_files_separates_excel_and_word(self):
        """Test file categorization by document type."""
        processor = BatchProcessor("CustomerDocument/")
        all_files = processor.find_all_files()
        categorized = processor._categorize_files(all_files)

        # Verify structure
        assert "excel" in categorized
        assert "word" in categorized

        # Verify Excel formats
        assert ".xlsx" in categorized["excel"]
        assert ".xls" in categorized["excel"]
        assert ".xlsm" in categorized["excel"]

        # Verify Word formats
        assert ".docx" in categorized["word"]
        assert ".doc" in categorized["word"]
        assert ".docm" in categorized["word"]

        # Count totals
        word_total = sum(len(files) for files in categorized["word"].values())
        excel_total = sum(len(files) for files in categorized["excel"].values())

        print(f"\nCategorization results:")
        print(f"  Excel: {excel_total} files")
        print(f"  Word: {word_total} files")

        assert word_total >= 51, f"Expected at least 51 Word files, found {word_total}"
        assert excel_total >= 58, f"Expected at least 58 Excel files, found {excel_total}"


class TestWordExtraction:
    """Test Word document extraction with semantic chunking."""

    @pytest.fixture
    def word_files(self) -> List[Path]:
        """Get list of Word files for testing."""
        processor = BatchProcessor("CustomerDocument/")
        all_files = processor.find_all_files()
        word_files = [f for f in all_files if f.suffix.lower() in [".docx", ".doc", ".docm"]]
        # Filter out preprocessed files
        word_files = [f for f in word_files if "_preprocessed" not in f.name]
        return word_files

    @pytest.mark.asyncio
    async def test_word_extraction_single_file(self, word_files):
        """Test extraction on a single Word file."""
        if not word_files:
            pytest.skip("No Word files found for testing")

        # Use first .docx file (most compatible format)
        docx_files = [f for f in word_files if f.suffix.lower() == ".docx"]
        if not docx_files:
            pytest.skip("No .docx files found for testing")

        test_file = docx_files[0]
        print(f"\nTesting single file: {test_file.name}")

        pipeline = ExtractionPipelineV2()
        doc_id = uuid4()

        result = await pipeline.process_word_document(doc_id, str(test_file))

        # Verify successful extraction
        assert result.success, f"Extraction failed: {result.error}"
        assert result.record_count > 0, "No chunks created"

        print(f"✓ Extraction successful:")
        print(f"  Chunks: {result.record_count}")
        print(f"  File: {test_file.name}")

        # Verify chunk structure
        assert len(result.records) == result.record_count
        first_chunk = result.records[0]
        assert "text" in first_chunk
        assert "metadata" in first_chunk
        assert "file_id" in first_chunk

    @pytest.mark.asyncio
    async def test_all_word_files_extraction(self, word_files):
        """Test extraction on all 51 Word files."""
        if not word_files:
            pytest.skip("No Word files found for testing")

        print(f"\nTesting all {len(word_files)} Word files...")
        pipeline = ExtractionPipelineV2()

        results = []
        for i, file_path in enumerate(word_files, 1):
            print(f"\n[{i}/{len(word_files)}] Processing: {file_path.name}")

            try:
                doc_id = uuid4()
                result = await pipeline.process_word_document(doc_id, str(file_path))

                results.append({
                    "file": file_path.name,
                    "format": file_path.suffix.lower(),
                    "success": result.success,
                    "chunks": result.record_count if result.success else 0,
                    "error": result.error if not result.success else None
                })

                if result.success:
                    print(f"  ✓ Success: {result.record_count} chunks")
                else:
                    print(f"  ✗ Failed: {result.error}")

            except Exception as e:
                print(f"  ✗ Exception: {e}")
                results.append({
                    "file": file_path.name,
                    "format": file_path.suffix.lower(),
                    "success": False,
                    "chunks": 0,
                    "error": str(e)
                })

        # Calculate success rate
        successful = sum(1 for r in results if r["success"])
        total = len(results)
        success_rate = (successful / total * 100) if total > 0 else 0

        # Print summary
        print(f"\n{'='*70}")
        print(f"WORD EXTRACTION TEST RESULTS")
        print(f"{'='*70}")
        print(f"Total files: {total}")
        print(f"  ✓ Success: {successful}")
        print(f"  ✗ Failed: {total - successful}")
        print(f"Success Rate: {success_rate:.2f}%")
        print()

        # Breakdown by format
        for ext in [".docx", ".doc", ".docm"]:
            ext_results = [r for r in results if r["format"] == ext]
            ext_success = sum(1 for r in ext_results if r["success"])
            ext_total = len(ext_results)
            ext_rate = (ext_success / ext_total * 100) if ext_total > 0 else 0
            print(f"{ext}: {ext_success}/{ext_total} ({ext_rate:.1f}%)")

        print(f"{'='*70}\n")

        # List failures
        if successful < total:
            print("Failed files:")
            for r in results:
                if not r["success"]:
                    print(f"  ✗ {r['file']}: {r['error']}")
            print()

        # Assert success rate meets target
        assert success_rate >= 85, (
            f"Success rate {success_rate:.2f}% below 85% target. "
            f"Failed {total - successful}/{total} files."
        )


class TestBatchProcessorMetrics:
    """Test batch processor metrics generation."""

    def test_generate_metrics_summary_structure(self):
        """Test that metrics summary has correct structure."""
        processor = BatchProcessor("CustomerDocument/")

        # Initialize with some files
        all_files = processor.find_all_files()
        processor.initialize_files(all_files)

        # Generate metrics
        metrics = processor.generate_metrics_summary()

        # Verify structure
        assert "summary" in metrics
        assert "by_type" in metrics
        assert "by_format" in metrics

        # Verify summary keys
        assert "total_files" in metrics["summary"]
        assert "completed" in metrics["summary"]
        assert "failed" in metrics["summary"]
        assert "success_rate" in metrics["summary"]

        # Verify by_type breakdown
        assert "excel" in metrics["by_type"]
        assert "word" in metrics["by_type"]

        # Verify Excel metrics
        assert "total" in metrics["by_type"]["excel"]
        assert "success" in metrics["by_type"]["excel"]
        assert "success_rate" in metrics["by_type"]["excel"]

        # Verify Word metrics
        assert "total" in metrics["by_type"]["word"]
        assert "success" in metrics["by_type"]["word"]
        assert "success_rate" in metrics["by_type"]["word"]

        # Verify by_format breakdown
        assert ".xlsx" in metrics["by_format"]
        assert ".xls" in metrics["by_format"]
        assert ".xlsm" in metrics["by_format"]
        assert ".docx" in metrics["by_format"]
        assert ".doc" in metrics["by_format"]
        assert ".docm" in metrics["by_format"]

        print(f"\n✓ Metrics structure valid")
        print(f"  Total files tracked: {metrics['summary']['total_files']}")
        print(f"  Excel files: {metrics['by_type']['excel']['total']}")
        print(f"  Word files: {metrics['by_type']['word']['total']}")

    def test_word_file_tracking_includes_chunk_metrics(self):
        """Test that Word files include chunk-specific metrics."""
        processor = BatchProcessor("CustomerDocument/")
        all_files = processor.find_all_files()
        processor.initialize_files(all_files)

        # Find a Word file in status
        word_file = None
        for file_key, file_info in processor.status_data["files"].items():
            if file_info.get("file_type") == "word":
                word_file = file_info
                break

        assert word_file is not None, "No Word files found in status"

        # Verify Word-specific fields exist
        assert "chunks" in word_file, "Word file missing 'chunks' field"
        assert "avg_chunk_size" in word_file, "Word file missing 'avg_chunk_size' field"
        assert "file_type" in word_file
        assert word_file["file_type"] == "word"

        print(f"\n✓ Word file metrics structure valid")
        print(f"  File: {word_file['filename']}")
        print(f"  Type: {word_file['file_type']}")
        print(f"  Format: {word_file['format']}")


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])
