"""Usage examples for VectorStore (Story 4.4).

This script demonstrates:
1. Creating Milvus collection
2. Embedding and indexing records
3. Similarity search
4. Integration with StructuredStore
5. Performance monitoring

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

import asyncio
from uuid import uuid4

from pymilvus import Collection

from src.core.config import settings
from src.extraction.models import ExtractedRecord
from src.knowledge.embeddings import EmbeddingService, record_to_text
from src.knowledge.vector_store import VectorStore


# ============================================================================
# Example 1: Text Representation
# ============================================================================


def example_1_text_representation():
    """Example 1: Convert records to text for embedding."""
    print("=" * 80)
    print("Example 1: Text Representation")
    print("=" * 80)

    # Simple record
    record1 = ExtractedRecord(
        content={
            "Item Code": "2024-4_0019",
            "Report": "not applicable",
            "Screen": "Applicable",
        },
        headers=["Item Code", "Report", "Screen"],
        _source={
            "document_id": str(uuid4()),
            "filename": "distribution_spec.xlsx",
            "sheet": "Item Mapping",
            "table_id": 1,
            "row": 42,
            "col_range": "A:D",
        },
    )

    text1 = record_to_text(record1)
    print(f"\nSimple Record:")
    print(f"  Content: {record1.content}")
    print(f"  Text:    {text1}")

    # Record with resolved symbols
    record2 = ExtractedRecord(
        content={"Item Code": "2024-4_0020", "Status": "○"},
        headers=["Item Code", "Status"],
        _source={
            "document_id": str(uuid4()),
            "filename": "distribution_spec.xlsx",
            "sheet": "Item Mapping",
            "table_id": 1,
            "row": 43,
            "col_range": "A:D",
        },
    )

    record2.resolved_content = {
        "Item Code": "2024-4_0020",
        "Status": "○",
        "Status_resolved": "Applicable/Yes",
    }

    text2 = record_to_text(record2)
    print(f"\nRecord with Symbols:")
    print(f"  Content:  {record2.content}")
    print(f"  Resolved: {record2.resolved_content}")
    print(f"  Text:     {text2}")
    print()


# ============================================================================
# Example 2: Embedding Service
# ============================================================================


async def example_2_embedding_service():
    """Example 2: Generate embeddings using Azure OpenAI."""
    print("=" * 80)
    print("Example 2: Embedding Service")
    print("=" * 80)

    if not settings.azure_openai_api_key:
        print("\n⚠️  Azure OpenAI credentials not configured")
        print("   Set AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT in .env")
        return

    embedding_service = EmbeddingService()

    texts = [
        "Item Code: 2024-4_0019. Report: not applicable. From: dist.xlsx, Sheet: Items, Row: 42",
        "Item Code: 2024-4_0020. Report: applicable. From: dist.xlsx, Sheet: Items, Row: 43",
    ]

    print(f"\nEmbedding {len(texts)} texts...")
    embeddings = await embedding_service.embed_batch(texts)

    print(f"\n✅ Generated {len(embeddings)} embeddings")
    print(f"   Model: {embedding_service.deployment}")
    print(f"   Dimensions: {len(embeddings[0])}")
    print(f"   First 5 values: {embeddings[0][:5]}")
    print()


# ============================================================================
# Example 3: Create Milvus Collection
# ============================================================================


async def example_3_create_collection():
    """Example 3: Create Milvus collection with schema."""
    print("=" * 80)
    print("Example 3: Create Milvus Collection")
    print("=" * 80)

    vector_store = VectorStore()

    print(f"\nCreating collection: {vector_store.collection_name}")
    print(f"  - Vector dimensions: {vector_store.vector_dim}")
    print(f"  - Index type: HNSW")
    print(f"  - Distance metric: COSINE")

    await vector_store.create_collection()

    print("\n✅ Collection created successfully")
    print()


# ============================================================================
# Example 4: Index Records
# ============================================================================


async def example_4_index_records():
    """Example 4: Embed and index records in Milvus."""
    print("=" * 80)
    print("Example 4: Index Records")
    print("=" * 80)

    if not settings.azure_openai_api_key:
        print("\n⚠️  Azure OpenAI credentials not configured")
        return

    # Create sample records
    document_id = uuid4()
    records = []

    for i in range(10):
        record = ExtractedRecord(
            content={
                "Item Code": f"TEST-{i:04d}",
                "Description": f"Test item {i}",
                "Status": "Active" if i % 2 == 0 else "Inactive",
            },
            headers=["Item Code", "Description", "Status"],
            _source={
                "document_id": str(document_id),
                "filename": "example_items.xlsx",
                "sheet": "Items",
                "table_id": 1,
                "row": i + 1,
                "col_range": "A:D",
            },
        )
        records.append(record)

    # Generate record IDs (normally from StructuredStore)
    record_ids = [uuid4() for _ in records]

    # Index records
    vector_store = VectorStore()
    await vector_store.create_collection()

    print(f"\nIndexing {len(records)} records...")
    count = await vector_store.index_records(records, document_id, record_ids)

    print(f"\n✅ Indexed {count} records")
    print()


# ============================================================================
# Example 5: Similarity Search
# ============================================================================


async def example_5_similarity_search():
    """Example 5: Perform semantic similarity search."""
    print("=" * 80)
    print("Example 5: Similarity Search")
    print("=" * 80)

    if not settings.azure_openai_api_key:
        print("\n⚠️  Azure OpenAI credentials not configured")
        return

    # First, index some records
    document_id = uuid4()
    records = [
        ExtractedRecord(
            content={
                "Item Code": "LAPTOP-001",
                "Description": "High performance laptop computer",
                "Category": "Electronics",
            },
            headers=["Item Code", "Description", "Category"],
            _source={
                "document_id": str(document_id),
                "filename": "products.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 1,
                "col_range": "A:C",
            },
        ),
        ExtractedRecord(
            content={
                "Item Code": "DESK-001",
                "Description": "Ergonomic standing desk with adjustable height",
                "Category": "Furniture",
            },
            headers=["Item Code", "Description", "Category"],
            _source={
                "document_id": str(document_id),
                "filename": "products.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 2,
                "col_range": "A:C",
            },
        ),
        ExtractedRecord(
            content={
                "Item Code": "MOUSE-001",
                "Description": "Wireless mouse with ergonomic design",
                "Category": "Electronics",
            },
            headers=["Item Code", "Description", "Category"],
            _source={
                "document_id": str(document_id),
                "filename": "products.xlsx",
                "sheet": "Products",
                "table_id": 1,
                "row": 3,
                "col_range": "A:C",
            },
        ),
    ]

    record_ids = [uuid4() for _ in records]

    # Index records
    vector_store = VectorStore()
    await vector_store.create_collection()
    await vector_store.index_records(records, document_id, record_ids)

    print("\nIndexed 3 product records")

    # Perform similarity search
    collection = Collection(settings.milvus_collection)
    collection.load()

    # Search for "computer accessories"
    embedding_service = EmbeddingService()
    query_text = "computer accessories for desk setup"
    query_embedding = await embedding_service.embed_batch([query_text])

    print(f"\nSearching for: '{query_text}'")

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

    print("\n✅ Search Results:")
    for i, hit in enumerate(search_results[0], 1):
        print(f"\n{i}. Similarity: {hit.distance:.4f}")
        print(f"   File: {hit.entity.get('filename')}")
        print(f"   Row:  {hit.entity.get('row_number')}")
        print(f"   Text: {hit.entity.get('text_content')}")

    print()


# ============================================================================
# Example 6: Integration with StructuredStore
# ============================================================================


async def example_6_full_integration():
    """Example 6: Complete workflow with StructuredStore and VectorStore."""
    print("=" * 80)
    print("Example 6: Full Integration Workflow")
    print("=" * 80)

    if not settings.azure_openai_api_key:
        print("\n⚠️  Azure OpenAI credentials not configured")
        return

    print("\nThis example shows the complete indexing workflow:")
    print("  1. Store records in PostgreSQL (StructuredStore)")
    print("  2. Embed and index in Milvus (VectorStore)")
    print("  3. Query both stores simultaneously")

    # Create sample records
    document_id = uuid4()
    records = [
        ExtractedRecord(
            content={"Item": "ITEM-001", "Value": "100"},
            headers=["Item", "Value"],
            _source={
                "document_id": str(document_id),
                "filename": "data.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": i,
                "col_range": "A:B",
            },
        )
        for i in range(1, 6)
    ]

    print(f"\nStep 1: Store {len(records)} records in PostgreSQL")
    print("        (Would use StructuredStore.store_records())")

    # Simulate record IDs from PostgreSQL
    record_ids = [uuid4() for _ in records]
    print(f"        Got {len(record_ids)} record IDs from database")

    print(f"\nStep 2: Embed and index in Milvus")
    vector_store = VectorStore()
    await vector_store.create_collection()
    count = await vector_store.index_records(records, document_id, record_ids)
    print(f"        ✅ Indexed {count} vectors in Milvus")

    print(f"\nStep 3: Both stores ready for querying")
    print(f"        - PostgreSQL: Exact lookups, JSONB queries")
    print(f"        - Milvus:     Semantic similarity search")

    print("\n✅ Dual-store indexing complete")
    print()


# ============================================================================
# Example 7: Performance Monitoring
# ============================================================================


async def example_7_performance_monitoring():
    """Example 7: Monitor indexing performance."""
    print("=" * 80)
    print("Example 7: Performance Monitoring")
    print("=" * 80)

    if not settings.azure_openai_api_key:
        print("\n⚠️  Azure OpenAI credentials not configured")
        return

    import time

    # Create test dataset
    document_id = uuid4()
    record_count = 50  # Scaled down for quick demo

    print(f"\nCreating {record_count} test records...")
    records = []
    for i in range(record_count):
        record = ExtractedRecord(
            content={
                "ID": f"REC-{i:05d}",
                "Data": f"Test record {i}",
                "Status": "Active",
            },
            headers=["ID", "Data", "Status"],
            _source={
                "document_id": str(document_id),
                "filename": "perf_test.xlsx",
                "sheet": "Data",
                "table_id": 1,
                "row": i + 1,
                "col_range": "A:C",
            },
        )
        records.append(record)

    record_ids = [uuid4() for _ in records]

    # Index and measure performance
    vector_store = VectorStore()
    await vector_store.create_collection()

    print(f"Indexing {record_count} records...")
    start_time = time.time()

    count = await vector_store.index_records(records, document_id, record_ids)

    duration = time.time() - start_time
    records_per_second = count / duration

    print(f"\n📊 Performance Metrics:")
    print(f"   Records indexed:     {count}")
    print(f"   Total duration:      {duration:.2f}s")
    print(f"   Records per second:  {records_per_second:.2f}")
    print(f"   Avg time per record: {(duration / count) * 1000:.2f}ms")

    # Project to 1000 records
    projected_time = 1000 / records_per_second
    print(f"\n📈 Projected Performance:")
    print(f"   Time for 1000 records: {projected_time:.2f}s")

    target = 30.0  # Target: ≤ 30s per 1000 records
    if projected_time <= target:
        print(f"   ✅ Within target ({target}s)")
    else:
        print(f"   ⚠️  Exceeds target ({target}s)")

    print()


# ============================================================================
# Example 8: Incremental Operations (Story 4.5)
# ============================================================================


async def example_8_incremental_operations():
    """Example 8: Delete and reindex operations."""
    print("=" * 80)
    print("Example 8: Incremental Operations (Story 4.5)")
    print("=" * 80)

    from src.db.session import get_async_session
    from src.knowledge.vector_store import delete_vectors_by_document_id
    from src.services.document_deletion_service import DocumentDeletionService

    document_id = uuid4()

    print("\n8a. Creating sample records...")
    sample_records = [
        ExtractedRecord(
            content={"Product": "Widget A", "Price": "100"},
            headers=["Product", "Price"],
            _source={
                "document_id": str(document_id),
                "filename": "products.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 2,
                "col_range": "A:B",
            },
        ),
        ExtractedRecord(
            content={"Product": "Widget B", "Price": "200"},
            headers=["Product", "Price"],
            _source={
                "document_id": str(document_id),
                "filename": "products.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 3,
                "col_range": "A:B",
            },
        ),
    ]

    print(f"  Document ID: {document_id}")
    print(f"  Records: {len(sample_records)}")

    print("\n8b. Deleting vectors by document ID...")
    try:
        vectors_deleted = await delete_vectors_by_document_id(document_id)
        print(f"  Vectors deleted: {vectors_deleted}")
    except Exception as e:
        print(f"  Error: {e}")

    print("\n8c. Document deletion service (atomic deletion)...")
    async with get_async_session() as session:
        deletion_service = DocumentDeletionService(session)

        try:
            # Note: This requires the document to exist in the database
            result = await deletion_service.delete_document(document_id)
            print(f"  Records deleted: {result['records']}")
            print(f"  Vectors deleted: {result['vectors']}")
        except Exception as e:
            print(f"  Error: {e}")
            print("  (This is expected if document doesn't exist)")

    print("\n8d. Reindex operation (delete old + index new)...")
    async with get_async_session() as session:
        deletion_service = DocumentDeletionService(session)

        # Updated records with different data
        updated_records = [
            ExtractedRecord(
                content={"Product": "Widget A", "Price": "150"},  # Price changed
                headers=["Product", "Price"],
                _source={
                    "document_id": str(document_id),
                    "filename": "products.xlsx",
                    "sheet": "Sheet1",
                    "table_id": 1,
                    "row": 2,
                    "col_range": "A:B",
                },
            ),
        ]

        try:
            result = await deletion_service.reindex_document(
                document_id, updated_records
            )
            print(f"  Deleted:")
            print(f"    Records: {result['deleted_records']}")
            print(f"    Vectors: {result['deleted_vectors']}")
            print(f"  Indexed:")
            print(f"    Records: {result['indexed_records']}")
            print(f"    Vectors: {result['indexed_vectors']}")
        except Exception as e:
            print(f"  Error: {e}")
            print("  (This is expected if document doesn't exist)")

    print()


# ============================================================================
# Main
# ============================================================================


async def main():
    """Run all examples."""
    print("\n" + "=" * 80)
    print("VectorStore Usage Examples (Stories 4.4 & 4.5)")
    print("=" * 80 + "\n")

    # Example 1: Text Representation (no API calls)
    example_1_text_representation()

    # Example 2: Embedding Service
    await example_2_embedding_service()

    # Example 3: Create Collection
    await example_3_create_collection()

    # Example 4: Index Records
    await example_4_index_records()

    # Example 5: Similarity Search
    await example_5_similarity_search()

    # Example 6: Full Integration
    await example_6_full_integration()

    # Example 7: Performance Monitoring
    await example_7_performance_monitoring()

    # Example 8: Incremental Operations (Story 4.5)
    await example_8_incremental_operations()

    print("=" * 80)
    print("All examples completed!")
    print("=" * 80 + "\n")


if __name__ == "__main__":
    asyncio.run(main())