"""Integration test for Story 4.5: Incremental Indexing.

This test demonstrates:
1. Indexing records in both PostgreSQL and Milvus
2. Deleting document and all associated records/vectors (atomic)
3. Re-indexing with updated records (delete old + index new)
4. Verifying data consistency between PostgreSQL and Milvus

Prerequisites:
- PostgreSQL running with database initialized
- Milvus running on localhost:19530
- Azure OpenAI credentials in .env
"""

import asyncio
from uuid import uuid4

from pymilvus import Collection, connections

from src.core.config import settings
from src.db.session import async_session_maker
from src.extraction.models import ExtractedRecord
from src.knowledge.structured_store import StructuredStore
from src.knowledge.vector_store import VectorStore, delete_vectors_by_document_id
from src.services.document_deletion_service import DocumentDeletionService


async def test_story_4_5():
    """Test Story 4.5: Incremental Indexing."""
    print("\n" + "=" * 80)
    print("Story 4.5: Incremental Indexing - Integration Test")
    print("=" * 80 + "\n")

    # Generate test document ID
    document_id = uuid4()
    print(f"Test Document ID: {document_id}\n")

    # Create document record (required for foreign key constraint)
    async with async_session_maker() as session:
        from src.db.models import Document

        document = Document(
            id=document_id,
            filename="inventory.xlsx",
            file_path=f"/tmp/test/{document_id}/inventory.xlsx",
            file_size_bytes=1024,
            status="completed",
        )
        session.add(document)
        await session.commit()
        print(f"✓ Created document record: {document_id}\n")

    # ========================================================================
    # Test 1: Index Initial Records (PostgreSQL + Milvus)
    # ========================================================================
    print("Test 1: Index Initial Records")
    print("-" * 80)

    initial_records = [
        ExtractedRecord(
            content={"Product": "Widget A", "Price": "100", "Stock": "50"},
            headers=["Product", "Price", "Stock"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 2,
                "col_range": "A:C",
            },
        ),
        ExtractedRecord(
            content={"Product": "Widget B", "Price": "200", "Stock": "30"},
            headers=["Product", "Price", "Stock"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 3,
                "col_range": "A:C",
            },
        ),
        ExtractedRecord(
            content={"Product": "Widget C", "Price": "150", "Stock": "20"},
            headers=["Product", "Price", "Stock"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 4,
                "col_range": "A:C",
            },
        ),
    ]

    async with async_session_maker() as session:
        # Index in PostgreSQL
        structured_store = StructuredStore(session)
        record_ids = await structured_store.store_records(initial_records, document_id)
        print(f"✓ PostgreSQL: Indexed {len(record_ids)} records")

        # Index in Milvus
        vector_store = VectorStore()
        await vector_store.create_collection()
        vector_count = await vector_store.index_records(
            initial_records, document_id, record_ids
        )
        print(f"✓ Milvus: Indexed {vector_count} vectors")

    # Verify records exist
    async with async_session_maker() as session:
        from sqlalchemy import select, func
        from src.db.models import ExtractedRecord as ExtractedRecordModel

        result = await session.execute(
            select(func.count(ExtractedRecordModel.id)).where(
                ExtractedRecordModel.document_id == document_id
            )
        )
        pg_count = result.scalar()
        print(f"✓ Verification: {pg_count} records in PostgreSQL")

    # Verify vectors exist in Milvus
    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )
    collection = Collection(settings.milvus_collection)
    collection.load()
    
    expr = f'document_id == "{str(document_id)}"'
    milvus_results = collection.query(expr=expr, output_fields=["id"])
    milvus_count = len(milvus_results)
    print(f"✓ Verification: {milvus_count} vectors in Milvus\n")

    connections.disconnect("default")

    # ========================================================================
    # Test 2: Delete Document (Atomic Deletion)
    # ========================================================================
    print("Test 2: Atomic Document Deletion")
    print("-" * 80)

    async with async_session_maker() as session:
        deletion_service = DocumentDeletionService(session)
        deletion_result = await deletion_service.delete_document(document_id)

        print(f"✓ Deleted {deletion_result['records']} records from PostgreSQL")
        print(f"✓ Deleted {deletion_result['vectors']} vectors from Milvus")

    # Verify deletion from PostgreSQL
    async with async_session_maker() as session:
        result = await session.execute(
            select(func.count(ExtractedRecordModel.id)).where(
                ExtractedRecordModel.document_id == document_id
            )
        )
        pg_count_after = result.scalar()
        print(f"✓ Verification: {pg_count_after} records remaining in PostgreSQL")

    # Verify deletion from Milvus
    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )
    collection = Collection(settings.milvus_collection)
    collection.load()
    
    milvus_results_after = collection.query(expr=expr, output_fields=["id"])
    milvus_count_after = len(milvus_results_after)
    print(f"✓ Verification: {milvus_count_after} vectors remaining in Milvus\n")

    connections.disconnect("default")

    # Assert deletion was complete
    assert pg_count_after == 0, f"Expected 0 PostgreSQL records, got {pg_count_after}"
    assert milvus_count_after == 0, f"Expected 0 Milvus vectors, got {milvus_count_after}"

    # ========================================================================
    # Test 3: Re-index Document (Delete Old + Index New)
    # ========================================================================
    print("Test 3: Re-index Document with Updated Records")
    print("-" * 80)

    # New records with updated data
    updated_records = [
        ExtractedRecord(
            content={"Product": "Widget A", "Price": "120", "Stock": "45"},  # Price/Stock changed
            headers=["Product", "Price", "Stock"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 2,
                "col_range": "A:C",
            },
        ),
        ExtractedRecord(
            content={"Product": "Widget B", "Price": "220", "Stock": "25"},  # Price/Stock changed
            headers=["Product", "Price", "Stock"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 3,
                "col_range": "A:C",
            },
        ),
        # Note: Widget C removed, Widget D added
        ExtractedRecord(
            content={"Product": "Widget D", "Price": "180", "Stock": "40"},  # New product
            headers=["Product", "Price", "Stock"],
            _source={
                "document_id": str(document_id),
                "filename": "inventory.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 4,
                "col_range": "A:C",
            },
        ),
    ]

    async with async_session_maker() as session:
        deletion_service = DocumentDeletionService(session)
        reindex_result = await deletion_service.reindex_document(
            document_id, updated_records
        )

        print(f"✓ Deleted {reindex_result['deleted_records']} old records")
        print(f"✓ Deleted {reindex_result['deleted_vectors']} old vectors")
        print(f"✓ Indexed {reindex_result['indexed_records']} new records")
        print(f"✓ Indexed {reindex_result['indexed_vectors']} new vectors")

    # Verify new records in PostgreSQL
    async with async_session_maker() as session:
        result = await session.execute(
            select(func.count(ExtractedRecordModel.id)).where(
                ExtractedRecordModel.document_id == document_id
            )
        )
        pg_count_final = result.scalar()
        print(f"✓ Verification: {pg_count_final} records in PostgreSQL after reindex")

        # Verify content was updated
        result = await session.execute(
            select(ExtractedRecordModel).where(
                ExtractedRecordModel.document_id == document_id
            ).order_by(ExtractedRecordModel.row_number)
        )
        records = result.scalars().all()
        
        print(f"✓ Record 1: {records[0].content}")  # Should show updated price/stock
        print(f"✓ Record 2: {records[1].content}")
        print(f"✓ Record 3: {records[2].content}")  # Should be Widget D

    # Verify new vectors in Milvus
    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )
    collection = Collection(settings.milvus_collection)
    collection.load()
    
    milvus_results_final = collection.query(expr=expr, output_fields=["id", "text"])
    milvus_count_final = len(milvus_results_final)
    print(f"✓ Verification: {milvus_count_final} vectors in Milvus after reindex\n")

    connections.disconnect("default")

    # Assert reindexing worked correctly
    assert pg_count_final == 3, f"Expected 3 PostgreSQL records, got {pg_count_final}"
    assert milvus_count_final == 3, f"Expected 3 Milvus vectors, got {milvus_count_final}"

    # ========================================================================
    # Test 4: Cleanup
    # ========================================================================
    print("Test 4: Final Cleanup")
    print("-" * 80)

    async with async_session_maker() as session:
        deletion_service = DocumentDeletionService(session)
        cleanup_result = await deletion_service.delete_document(document_id)

        print(f"✓ Cleaned up {cleanup_result['records']} records")
        print(f"✓ Cleaned up {cleanup_result['vectors']} vectors")

        # Also delete the document record
        from sqlalchemy import delete
        from src.db.models import Document

        await session.execute(delete(Document).where(Document.id == document_id))
        await session.commit()
        print(f"✓ Deleted document record\n")

    # ========================================================================
    # Summary
    # ========================================================================
    print("=" * 80)
    print("✅ Story 4.5 Integration Test PASSED")
    print("=" * 80)
    print("\nAll Acceptance Criteria Verified:")
    print("  ✓ AC4.5.1: Add new document records incrementally")
    print("  ✓ AC4.5.2: Delete document and all its records atomically")
    print("  ✓ AC4.5.3: Re-process document (delete + re-index)")
    print("  ✓ AC4.5.4: Integration with document lifecycle")
    print("  ✓ AC4.5.5: Error handling and rollback (tested via unit tests)")
    print("\nData Consistency:")
    print("  ✓ PostgreSQL and Milvus stay in sync")
    print("  ✓ Atomic operations prevent partial state")
    print("  ✓ Reindexing correctly updates both stores")
    print()


if __name__ == "__main__":
    try:
        asyncio.run(test_story_4_5())
    except Exception as e:
        print(f"\n❌ Test Failed: {e}")
        import traceback
        traceback.print_exc()
        exit(1)