#!/usr/bin/env python3
"""Manual test script for the multi-pass extraction pipeline.

This script demonstrates how to use the ExtractionPipeline directly
to process an Excel document and extract structured records with:
    PASS 1: Symbol Detection
    PASS 2: Table Extraction
    PASS 3: Symbol Resolution
    PASS 4: Cross-Reference Detection
"""

from pathlib import Path
import redis as sync_redis

from src.db.session import get_sync_session
from src.db import models as db_models
from src.extraction.pipeline import ExtractionPipeline  # V1 pipeline
from src.extraction_v2.pipeline import ExtractionPipelineV2  # V2 pipeline
from src.core.config import settings


def main():
    """Run the extraction pipeline on a test file."""

    # Choose a test file (you can change this to any Excel file)
    test_files = [
        # Test fixtures
        ("tests/fixtures/sample_excel/test_multi_table_crossref.xlsx", "Multi-table + cross-refs (recommended)"),
        ("tests/fixtures/sample_excel/test_multipass_comprehensive.xlsx", "Multi-pass test (symbols + cross-refs)"),
        ("tests/fixtures/sample_excel/test_standard_table.xlsx", "Standard table (simple)"),
        ("tests/fixtures/sample_excel/test_multiple_tables.xlsx", "Multiple tables"),
        ("tests/fixtures/sample_excel/test_mixed_depth_headers.xlsx", "Mixed depth headers"),
        ("tests/fixtures/sample_excel/test_complex_merges.xlsx", "Complex merged cells"),
        # Customer documents (paths from filesystem)
        ("CustomerDocument/IFф╗ХцзШцЫ╕(BB-CASTAR)/focus_1.xlsx", "Customer: BB-CASTAR"),
        ("CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.2 JP1шинхоЪщаЕчЫо.xlsx", "Customer: JP1 Settings"),
        ("CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.1 JP1уВ╣уВ▒уВ╖уВЩуГеуГ╝уГлф╕Ашжз.xlsx", "Customer: JP1 Schedule"),
        ("CustomerDocument/чФ╗щЭвшиншиИч╖и/02_чФ╗щЭвф╗ХцзШ/01_щА▓цНЧчобчРЖцйЯшГ╜/хИеч┤Щ_01.02.01.03.05_х╖еф║ЛщА▓цНЧщБЕх╗╢уГБуВзуГГуВп.xlsx", "Customer: Screen Design"),
        ("CustomerDocument/ц╡БщАЪщаЕчЫоф╗ХцзШцЫ╕/01.01.01_ц╡БщАЪщаЕчЫоф╗ХцзШцЫ╕уАРуВвуГ│уГПуВЩуГ│уГИуВЩуГлцКЬч▓ЛчЙИуАС.xlsx", "Customer: Distribution Spec"),
        ("CustomerDocument/IFф╗ХцзШцЫ╕(уВвуГ│уГПуВЩуГ│уГИуВЩуГл)/ч╡▒хРИHHC/01.03.03-06_SOAP FaultуВдуГ│уВ┐уГХуВзуГ╝уВ╣щЫ╗цЦЗ.xlsx", "Customer: SOAP Fault IF"),
        ('/Users/nguyen/Work/DSOL/CustomerDocument/IFф╗ХцзШцЫ╕(уВвуГ│уГПуВЩуГ│уГИуВЩуГл)/уВкуГ╝уВ┐уВЩхИ╢х╛б/1 хЕЙуВвуГ│уГПуВЩуГ│уГИуВЩуГлценхЛЩцФпцП┤уВ╖уВ╣уГЖуГая╝ИхИЖцЮРч│╗я╝Й/1.12 уГХуВбуВдуГлчХкхП╖уБиуГлуГ╝уГИуВ┐уВпуВЩхп╛х┐Ьшби.xls', "Customer: Complex Document"),
    ]

    print("=" * 70)
    print("MULTI-PASS EXTRACTION PIPELINE TEST")
    print("=" * 70)
    print("\nAvailable test files:")
    for i, (file, desc) in enumerate(test_files, 1):
        exists = "[OK]" if Path(file).exists() else "[--]"
        print(f"  {i}. {exists} {desc}")
        print(f"           {file}")

    choice = input(f"\nSelect a file (1-{len(test_files)}) or enter a custom path: ").strip()

    try:
        choice_num = int(choice)
        if 1 <= choice_num <= len(test_files):
            test_file = Path(test_files[choice_num - 1][0])
        else:
            print("Invalid choice, using default file")
            test_file = Path(test_files[0][0])
    except ValueError:
        # User entered a custom path
        test_file = Path(choice)

    if not test_file.exists():
        print(f"Error: File not found: {test_file}")
        return

    print(f"\nUsing file: {test_file}")
    print("=" * 70)

    # Select pipeline version
    print("\nSelect extraction pipeline:")
    print("  1. V1 Pipeline (openpyxl-based, original)")
    print("  2. V2 Pipeline (Docling-based, with LLM headers)")
    print("  3. Both (compare side-by-side)")

    pipeline_choice = input("Pipeline (1-3, default=2): ").strip() or "2"

    if pipeline_choice not in ["1", "2", "3"]:
        print("Invalid choice, using V2 pipeline")
        pipeline_choice = "2"

    print("=" * 70)

    # Create a test document in database
    with get_sync_session() as session:
        # Create test document
        doc = db_models.Document(
            filename=test_file.name,
            file_path=str(test_file.absolute()),
            status="pending"
        )
        session.add(doc)
        session.commit()

        doc_id = doc.id
        print(f"\n[OK] Created test document: {doc_id}")

        # Optional: Connect to Redis for progress tracking
        redis_client = None
        try:
            redis_client = sync_redis.from_url(settings.redis_url, decode_responses=False)
            print("[OK] Redis connected (progress tracking enabled)")
        except Exception as e:
            print(f"[WARN] Redis not available: {e}")
            print("  (Proceeding without progress tracking)")

        print(f"\n[START] Running extraction pipeline(s)...")
        print("=" * 70)

        try:
            v1_summary = None
            v2_result = None

            # Run V1 pipeline
            if pipeline_choice in ["1", "3"]:
                print("\n[V1] Running openpyxl-based extraction pipeline...")
                print("-" * 70)
                v1_pipeline = ExtractionPipeline(db_session=session, redis_client=redis_client)
                v1_summary = v1_pipeline.process_document(doc_id)
                print(f"[V1] ✓ Complete: {v1_summary.records_extracted} records extracted")

            # Run V2 pipeline
            if pipeline_choice in ["2", "3"]:
                print("\n[V2] Running Docling-based extraction pipeline...")
                print("-" * 70)
                v2_pipeline = ExtractionPipelineV2()
                v2_result = v2_pipeline.process_document_with_large_table_detection_sync(doc_id, str(test_file.absolute()))

                if v2_result.success:
                    print(f"[V2] ✓ Complete: {v2_result.record_count} records extracted")

                    # Save V2 records to database
                    from src.workers.extraction_v2_tasks import _save_v2_records, _mark_document_completed, _index_vectors_sync
                    _save_v2_records(session, v2_result)
                    _mark_document_completed(session, str(doc_id), v2_result)

                    # Index vectors in Milvus
                    print("[V2] Indexing vectors in Milvus...")
                    try:
                        vector_count = _index_vectors_sync(v2_result.records, doc_id, str(test_file.absolute()))
                        print(f"[V2] ✓ Indexed {vector_count} vectors in Milvus")
                    except Exception as e:
                        print(f"[V2] ✗ Vector indexing failed: {e}")
                else:
                    print(f"[V2] ✗ Failed: {v2_result.error}")
                    print("[V2] Note: In production, this would trigger V1 fallback")

            # Print results
            print("\n" + "=" * 70)
            print("EXTRACTION COMPLETE!")
            print("=" * 70)

            # Show comparison if both were run
            if pipeline_choice == "3" and v1_summary and v2_result and v2_result.success:
                print("\n--- V1 vs V2 Comparison ---")
                print(f"  Records extracted:")
                print(f"    V1 (openpyxl):  {v1_summary.records_extracted}")
                print(f"    V2 (Docling):   {v2_result.record_count}")
                print(f"  Tables detected:")
                print(f"    V1: {v1_summary.tables_detected}")
                print(f"    V2: {v2_result.table_count}")
                print(f"  Processing time:")
                print(f"    V1: {v1_summary.duration_ms}ms ({v1_summary.duration_ms/1000:.2f}s)")
                print(f"    V2: (not tracked separately)")

                # Determine which to show detailed results for
                summary = v1_summary  # Default to V1 for detailed output
                print("\n  Showing detailed V1 results below (V2 records also saved to DB)")
            elif v2_result and v2_result.success:
                # V2 only - create a summary-like object for display
                summary = type('obj', (object,), {
                    'symbols_detected': 0,
                    'sheets_processed': v2_result.table_count,
                    'tables_detected': v2_result.table_count,
                    'records_extracted': v2_result.record_count,
                    'symbols_resolved': 0,
                    'unresolved_symbols': 0,
                    'cross_references_found': 0,
                    'cross_references_resolved': 0,
                    'cross_references_unresolved': 0,
                    'duration_ms': 0,
                    'warnings': []
                })()
            else:
                # V1 only or V2 failed
                summary = v1_summary

            # Display detailed results
            if pipeline_choice == "1" or (pipeline_choice == "3" and v1_summary):
                # Full V1 multi-pass results
                print("\n--- PASS 1: Symbol Detection (V1 only) ---")
                print(f"  Symbols detected: {summary.symbols_detected}")

                print("\n--- PASS 2: Table Extraction ---")
                print(f"  Sheets processed: {summary.sheets_processed}")
                print(f"  Tables detected: {summary.tables_detected}")
                print(f"  Records extracted: {summary.records_extracted}")

                print("\n--- PASS 3: Symbol Resolution (V1 only) ---")
                print(f"  Symbols resolved: {summary.symbols_resolved}")
                print(f"  Unresolved symbols: {summary.unresolved_symbols}")

                print("\n--- PASS 4: Cross-Reference Detection (V1 only) ---")
                print(f"  Cross-references found: {summary.cross_references_found}")
                print(f"  Cross-references resolved: {summary.cross_references_resolved}")
                print(f"  Cross-references unresolved: {summary.cross_references_unresolved}")

                print("\n--- Performance ---")
                print(f"  Duration: {summary.duration_ms}ms ({summary.duration_ms/1000:.2f}s)")
            else:
                # V2 results (simpler - just table extraction)
                print("\n--- V2 Extraction Results ---")
                print(f"  Tables extracted: {summary.tables_detected}")
                print(f"  Records extracted: {summary.records_extracted}")
                print("  Note: V2 pipeline focuses on table extraction with LLM-based headers")
                print("        Symbol resolution and cross-refs would run in subsequent passes")

            if summary.warnings:
                print(f"\n--- Warnings ({len(summary.warnings)}) ---")
                for i, warning in enumerate(summary.warnings, 1):
                    print(f"  {i}. {warning}")

            # Check database
            print(f"\n--- Database Status ---")
            print(f"  Document status: {doc.status}")
            print(f"  Record count: {doc.record_count}")
            print(f"  Sheet count: {doc.sheet_count}")
            print(f"  Processed at: {doc.processed_at}")

            # Query symbol dictionaries
            symbols = session.query(db_models.SymbolDictionary).filter_by(
                document_id=doc_id
            ).all()
            print(f"\n--- Symbol Dictionaries in Database ---")
            print(f"  Total symbols: {len(symbols)}")
            if symbols:
                print(f"  All symbols:")
                for i, sym in enumerate(symbols, 1):
                    print(f"    {i}. '{sym.symbol}' = '{sym.meaning}' (sheet: {sym.sheet_name})")

            # Query extracted records
            records = session.query(db_models.ExtractedRecord).filter_by(
                document_id=doc_id
            ).order_by(db_models.ExtractedRecord.sheet_name, db_models.ExtractedRecord.row_number).all()
            print(f"\n--- Extracted Records in Database ---")
            print(f"  Total records: {len(records)}")

            if records:
                # Group by sheet
                sheets = {}
                for record in records:
                    if record.sheet_name not in sheets:
                        sheets[record.sheet_name] = []
                    sheets[record.sheet_name].append(record)

                print(f"  Records by sheet:")
                for sheet_name, sheet_records in sheets.items():
                    print(f"    - {sheet_name}: {len(sheet_records)} records")

                # Show all records
                print(f"\n--- All Extracted Records ---")
                for i, record in enumerate(records, 1):
                    title_str = f", Table: {record.table_title}" if record.table_title else ""
                    print(f"\n  Record {i} (Sheet: {record.sheet_name}{title_str}, Row: {record.row_number}):")
                    print(f"    Headers: {record.headers}")
                    print(f"    Content: {record.content}")
                    if record.resolved_content:
                        print(f"    Resolved: {record.resolved_content}")

            # Query cross-references (through source_record relationship)
            cross_refs = session.query(db_models.CrossReference).join(
                db_models.ExtractedRecord,
                db_models.CrossReference.source_record_id == db_models.ExtractedRecord.id
            ).filter(
                db_models.ExtractedRecord.document_id == doc_id
            ).all()
            print(f"\n--- Cross-References in Database ---")
            print(f"  Total cross-references: {len(cross_refs)}")
            if cross_refs:
                resolved = [cr for cr in cross_refs if cr.resolved]
                unresolved = [cr for cr in cross_refs if not cr.resolved]
                print(f"  Resolved: {len(resolved)}")
                print(f"  Unresolved: {len(unresolved)}")

                print(f"\n  All cross-references:")
                for i, cr in enumerate(cross_refs, 1):
                    status = "[RESOLVED]" if cr.resolved else "[UNRESOLVED]"
                    print(f"    {i}. {status} '{cr.reference_text}'")
                    print(f"       Type: {cr.reference_type}")

            # Show Redis progress (if available)
            if redis_client:
                try:
                    key = f"processing:documents:{doc_id}"
                    progress = redis_client.get(key)
                    if progress:
                        import json
                        progress_data = json.loads(progress)
                        print(f"\n--- Redis Progress Data ---")
                        print(f"  {progress_data}")
                    else:
                        print(f"\n--- Redis: No progress data (expired or cleared) ---")
                except Exception as e:
                    print(f"\n--- Redis: Could not fetch progress: {e} ---")

            print("\n" + "=" * 70)
            print("Test completed successfully!")
            print("=" * 70)

        except Exception as e:
            print("\n" + "=" * 70)
            print("EXTRACTION FAILED!")
            print("=" * 70)
            print(f"\nError: {e}")
            print(f"\nDocument status: {doc.status}")
            if doc.error_message:
                print(f"Error message: {doc.error_message}")

            import traceback
            print("\nFull traceback:")
            traceback.print_exc()


if __name__ == "__main__":
    main()
