"""Integration test for Stories 4.1 through 4.4.

Tests the complete knowledge indexing workflow:
- Story 4.1: Store records in PostgreSQL (StructuredStore)
- Story 4.2: Verify source metadata preservation
- Story 4.3: Test exact lookup indexing
- Story 4.4: Embed and index in Milvus (VectorStore)

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

Usage:
    python examples/test_stories_4_1_to_4_4.py
"""

import asyncio
import sys
from uuid import uuid4

from pymilvus import Collection, connections
from sqlalchemy import select

from src.core.config import settings
from src.db.models import Document, ExtractedRecord as DBExtractedRecord
from src.db.session import async_session_maker
from src.extraction.models import ExtractedRecord
from src.knowledge.embeddings import record_to_text
from src.knowledge.indexer import Indexer
from src.knowledge.metadata_service import MetadataService
from src.knowledge.structured_store import StructuredStore
from src.knowledge.vector_store import VectorStore


# ============================================================================
# Mock Data Creation
# ============================================================================


def create_mock_records(document_id: str, filename: str) -> list[ExtractedRecord]:
    """Create mock ExtractedRecord data for testing.

    Simulates records from a distribution specification document with:
    - Regular content fields
    - Symbol fields (that could be resolved)
    - Various data types (codes, descriptions, status)
    """
    records = [
        # Record 1: Item with applicable report and screen
        ExtractedRecord(
            content={
                "Item Code": "2024-4_0019",
                "Item Name": "Distribution Module A",
                "Report": "not applicable",
                "Screen": "Applicable",
                "Notes": "Standard distribution item",
            },
            headers=["Item Code", "Item Name", "Report", "Screen", "Notes"],
            _source={
                "document_id": document_id,
                "filename": filename,
                "sheet": "Item Mapping",
                "table_id": 1,
                "row": 10,
                "col_range": "A:E",
            },
        ),

        # Record 2: Item with symbol that could be resolved
        ExtractedRecord(
            content={
                "Item Code": "2024-4_0020",
                "Item Name": "Distribution Module B",
                "Report": "applicable",
                "Screen": "○",  # Symbol
                "Notes": "Requires special handling",
            },
            headers=["Item Code", "Item Name", "Report", "Screen", "Notes"],
            _source={
                "document_id": document_id,
                "filename": filename,
                "sheet": "Item Mapping",
                "table_id": 1,
                "row": 11,
                "col_range": "A:E",
            },
        ),

        # Record 3: Another item with different status
        ExtractedRecord(
            content={
                "Item Code": "2024-4_0021",
                "Item Name": "Distribution Module C",
                "Report": "applicable",
                "Screen": "Not Applicable",
                "Notes": "Legacy system compatibility",
            },
            headers=["Item Code", "Item Name", "Report", "Screen", "Notes"],
            _source={
                "document_id": document_id,
                "filename": filename,
                "sheet": "Item Mapping",
                "table_id": 1,
                "row": 12,
                "col_range": "A:E",
            },
        ),

        # Record 4: Hardware item
        ExtractedRecord(
            content={
                "Item Code": "HW-2024-001",
                "Item Name": "Server Hardware Package",
                "Report": "applicable",
                "Screen": "Applicable",
                "Notes": "Includes rack mounting equipment",
            },
            headers=["Item Code", "Item Name", "Report", "Screen", "Notes"],
            _source={
                "document_id": document_id,
                "filename": filename,
                "sheet": "Hardware Items",
                "table_id": 2,
                "row": 5,
                "col_range": "A:E",
            },
        ),

        # Record 5: Software item
        ExtractedRecord(
            content={
                "Item Code": "SW-2024-001",
                "Item Name": "Database Management System",
                "Report": "not applicable",
                "Screen": "×",  # Symbol for "not applicable"
                "Notes": "PostgreSQL 15 with extensions",
            },
            headers=["Item Code", "Item Name", "Report", "Screen", "Notes"],
            _source={
                "document_id": document_id,
                "filename": filename,
                "sheet": "Software Items",
                "table_id": 3,
                "row": 8,
                "col_range": "A:E",
            },
        ),
    ]

    # Add resolved_content to records with symbols
    records[1].resolved_content = {
        "Item Code": "2024-4_0020",
        "Item Name": "Distribution Module B",
        "Report": "applicable",
        "Screen": "○",
        "Screen_resolved": "Applicable/Yes",
        "Notes": "Requires special handling",
    }

    records[4].resolved_content = {
        "Item Code": "SW-2024-001",
        "Item Name": "Database Management System",
        "Report": "not applicable",
        "Screen": "×",
        "Screen_resolved": "Not Applicable/No",
        "Notes": "PostgreSQL 15 with extensions",
    }

    return records


# ============================================================================
# Test Story 4.1: Structured Store (PostgreSQL)
# ============================================================================


async def test_story_4_1_structured_store(
    session,
    document_id: str,
    records: list[ExtractedRecord]
) -> list[str]:
    """Test Story 4.1: Store records in PostgreSQL."""
    print("\n" + "=" * 80)
    print("Story 4.1: Structured Store (PostgreSQL)")
    print("=" * 80)

    structured_store = StructuredStore(session)

    print(f"\nStoring {len(records)} records in PostgreSQL...")
    count = await structured_store.store_records(records, document_id)

    print(f"✅ Stored {count} records")

    # Query to get the actual record IDs
    result = await session.execute(
        select(DBExtractedRecord).where(
            DBExtractedRecord.document_id == document_id
        )
    )
    db_records = result.scalars().all()
    record_ids = [str(record.id) for record in db_records]

    print(f"   Sample record ID: {record_ids[0]}")

    print(f"\n✅ Verification:")
    print(f"   Records in database: {len(db_records)}")
    print(f"   Content stored: {len(db_records[0].content)} fields")
    print(f"   Headers stored: {len(db_records[0].headers)}")

    # Show sample record (db_records are SQLAlchemy models, not extraction models)
    sample = db_records[0]
    print(f"\n📄 Sample Record:")
    print(f"   ID: {sample.id}")
    print(f"   Document ID: {sample.document_id}")
    print(f"   Content: {sample.content}")
    print(f"   Headers: {sample.headers}")

    return record_ids, db_records


# ============================================================================
# Test Story 4.2: Source Metadata Preservation
# ============================================================================


async def test_story_4_2_metadata_service(
    session,
    document_id: str,
    db_records: list
):
    """Test Story 4.2: Verify source metadata preservation."""
    print("\n" + "=" * 80)
    print("Story 4.2: Source Metadata Preservation")
    print("=" * 80)

    metadata_service = MetadataService(session)

    # Get source citation for first record
    first_record_id = db_records[0].id
    citation = await metadata_service.get_source_citation(first_record_id)
    if citation:
        print(f"\n✅ Retrieved source citation:")
        print(f"   Filename: {citation.filename}")
        print(f"   Sheet: {citation.sheet_name}")
        print(f"   Row: {citation.row_number}")

    # Get records by location
    records_by_location = await metadata_service.get_records_by_location(
        document_id, "Item Mapping"
    )
    print(f"✅ Found {len(records_by_location)} records in 'Item Mapping' sheet")

    # Validate source completeness
    validation = await metadata_service.validate_source_completeness(document_id)
    print(f"\n✅ Metadata validation:")
    print(f"   Total records: {validation.total_records}")
    print(f"   Complete: {validation.complete_records}")
    print(f"   Valid: {validation.is_valid}")
    print(f"   Completeness: {validation.completeness_percentage:.1f}%")


# ============================================================================
# Test Story 4.3: Exact Lookup Indexing
# ============================================================================


async def test_story_4_3_indexing(session, document_id: str):
    """Test Story 4.3: Verify indexes and test exact lookups."""
    print("\n" + "=" * 80)
    print("Story 4.3: Exact Lookup Indexing")
    print("=" * 80)

    indexer = Indexer(session)

    # Verify indexes exist
    print(f"\n🔍 Verifying database indexes...")
    index_report = await indexer.verify_indexes()

    if index_report.all_indexes_exist:
        print(f"   ✅ All {len(index_report.expected_indexes)} indexes exist")
    else:
        print(f"   ⚠️  Missing indexes: {len(index_report.missing_indexes)}")
        for idx in index_report.missing_indexes:
            print(f"      - {idx}")

    # Show index details
    print(f"\n📊 Index Details:")
    for idx_info in index_report.expected_indexes[:3]:
        status = "✅" if idx_info.exists else "❌"
        print(f"   {status} {idx_info.index_name}")
        print(f"      Table: {idx_info.table_name}")
        print(f"      Type: {idx_info.index_type}")

    # Get index statistics
    print(f"\n📈 Index Usage Statistics:")
    stats = await indexer.get_index_stats()

    for stat in stats[:3]:
        print(f"   {stat.index_name}:")
        print(f"      Scans: {stat.scans}")
        print(f"      Tuples read: {stat.tuples_read}")

    # Benchmark query performance
    print(f"\n⚡ Benchmarking Query Performance:")
    perf_report = await indexer.benchmark_query_performance(document_id)

    queries = [
        ("Exact Lookup", perf_report.exact_lookup_ms, 50),
        ("Symbol Lookup", perf_report.symbol_lookup_ms, 20),
        ("JSONB Containment", perf_report.jsonb_containment_ms, 100),
        ("Resolved Search", perf_report.resolved_search_ms, 200),
    ]

    for query_name, duration, target in queries:
        status = "✅" if duration < target else "⚠️"
        print(f"   {status} {query_name}: {duration:.2f}ms (target: <{target}ms)")


# ============================================================================
# Test Story 4.4: Vector Embedding and Indexing
# ============================================================================


async def test_story_4_4_vector_store(
    document_id: str,
    records: list[ExtractedRecord],
    record_ids: list[str]
):
    """Test Story 4.4: Embed and index in Milvus."""
    print("\n" + "=" * 80)
    print("Story 4.4: Vector Embedding and Indexing")
    print("=" * 80)

    if not settings.azure_openai_api_key:
        print("\n⚠️  Azure OpenAI credentials not configured")
        print("   Set AZURE_OPENAI_API_KEY in .env to test Story 4.4")
        return

    vector_store = VectorStore()

    # Create Milvus collection
    print(f"\n📦 Creating Milvus collection...")
    await vector_store.create_collection()
    print(f"   ✅ Collection: {vector_store.collection_name}")
    print(f"   ✅ Vector dimensions: {vector_store.vector_dim}")

    # Show text representations
    print(f"\n📝 Text Representations:")
    for i, record in enumerate(records[:3], 1):
        text = record_to_text(record)
        print(f"   {i}. {text[:80]}...")

    # Index records in Milvus
    print(f"\n🚀 Indexing {len(records)} records in Milvus...")
    count = await vector_store.index_records(records, document_id, record_ids)
    print(f"   ✅ Indexed {count} vectors")

    # Verify vectors in Milvus
    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )

    collection = Collection(settings.milvus_collection)
    collection.load()

    result = collection.query(
        expr=f'document_id == "{str(document_id)}"',
        output_fields=["record_id", "filename", "sheet_name", "text_content"],
    )

    print(f"\n✅ Verification:")
    print(f"   Vectors in Milvus: {len(result)}")
    if result:
        print(f"   Sample metadata: {result[0]['filename']}, {result[0]['sheet_name']}")
    else:
        print(f"   ⚠️ No vectors found (may need flush/load)")

    # Perform similarity search
    print(f"\n🔍 Testing Similarity Search:")

    from src.knowledge.embeddings import EmbeddingService
    embedding_service = EmbeddingService()

    queries = [
        "distribution modules with applicable reports",
        "hardware server equipment",
        "database software systems",
    ]

    for query_text in queries:
        print(f"\n   Query: '{query_text}'")

        query_embedding = await embedding_service.embed_batch([query_text])

        search_results = collection.search(
            data=query_embedding,
            anns_field="vector",
            param={"metric_type": "COSINE", "params": {"ef": 64}},
            limit=3,
            output_fields=["text_content", "sheet_name", "row_number"],
        )

        if search_results and search_results[0]:
            for i, hit in enumerate(search_results[0], 1):
                print(f"      {i}. Score: {hit.distance:.4f}")
                print(f"         {hit.entity.get('sheet_name')} (Row {hit.entity.get('row_number')})")
                print(f"         {hit.entity.get('text_content')[:60]}...")
        else:
            print(f"      No results found")

    connections.disconnect("default")


# ============================================================================
# Test Complete Workflow
# ============================================================================


async def test_dual_store_workflow(
    session,
    document_id: str,
    records: list[ExtractedRecord]
):
    """Test complete dual-store workflow."""
    print("\n" + "=" * 80)
    print("Complete Dual-Store Workflow Test")
    print("=" * 80)

    print(f"\n📊 Workflow Summary:")
    print(f"   Document ID: {document_id}")
    print(f"   Records: {len(records)}")

    # Count records by sheet
    sheets = {}
    for record in records:
        sheet = record._source["sheet"]
        sheets[sheet] = sheets.get(sheet, 0) + 1

    print(f"   Sheets:")
    for sheet, count in sheets.items():
        print(f"      - {sheet}: {count} records")

    print(f"\n💾 Storage Strategy:")
    print(f"   PostgreSQL (Story 4.1):")
    print(f"      - Structured data with JSONB")
    print(f"      - Exact lookups and filters")
    print(f"      - Source metadata preservation")
    print(f"      - B-tree and GIN indexes")

    print(f"\n   Milvus (Story 4.4):")
    print(f"      - 3072-dimensional vectors")
    print(f"      - Semantic similarity search")
    print(f"      - HNSW index for fast retrieval")
    print(f"      - Cosine distance metric")

    print(f"\n🎯 Use Cases:")
    print(f"   Exact Queries → PostgreSQL:")
    print(f"      - Find by Item Code: 'SW-2024-001'")
    print(f"      - Filter by sheet: 'Hardware Items'")
    print(f"      - JSONB queries: Report = 'applicable'")

    print(f"\n   Semantic Queries → Milvus:")
    print(f"      - 'What items require special handling?'")
    print(f"      - 'Find server and hardware packages'")
    print(f"      - 'Show database-related items'")

    print(f"\n✅ Both stores are now ready for querying!")


# ============================================================================
# Main Test Runner
# ============================================================================


async def main():
    """Run complete integration test for Stories 4.1-4.4."""
    print("\n" + "=" * 80)
    print("INTEGRATION TEST: Stories 4.1 through 4.4")
    print("Knowledge Base & Indexing")
    print("=" * 80)

    # Create mock document
    document_id = uuid4()
    filename = "distribution_spec_2024.xlsx"

    # Create mock records
    print(f"\n📝 Creating mock data...")
    records = create_mock_records(str(document_id), filename)
    print(f"   ✅ Created {len(records)} mock records")

    # Test each story
    async with async_session_maker() as session:
        # Create document entry
        document = Document(
            id=document_id,
            filename=filename,
            file_path=f"/uploads/{document_id}.xlsx",
            file_size_bytes=1024 * 500,  # 500KB
            status="completed",
        )
        session.add(document)
        await session.commit()

        print(f"   ✅ Created document: {filename}")

        # Story 4.1: Structured Store
        record_ids, db_records = await test_story_4_1_structured_store(
            session, str(document_id), records
        )

        # Story 4.2: Metadata Service
        await test_story_4_2_metadata_service(session, document_id, db_records)

        # Story 4.3: Indexing
        await test_story_4_3_indexing(session, document_id)

        # Story 4.4: Vector Store
        await test_story_4_4_vector_store(document_id, records, record_ids)

        # Complete workflow summary
        await test_dual_store_workflow(session, document_id, records)

    print("\n" + "=" * 80)
    print("✅ ALL TESTS COMPLETED SUCCESSFULLY")
    print("=" * 80)
    print(f"\nStories tested:")
    print(f"  ✅ Story 4.1: Structured Store (PostgreSQL)")
    print(f"  ✅ Story 4.2: Source Metadata Preservation")
    print(f"  ✅ Story 4.3: Exact Lookup Indexing")
    print(f"  ✅ Story 4.4: Vector Embedding and Indexing")
    print(f"\n🎉 Knowledge indexing workflow is fully operational!")
    print("=" * 80 + "\n")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n\n⚠️  Test interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)