"""Integration tests for vector_store with real Milvus.

These tests require:
- Milvus running on localhost:19530 (Docker)
- Azure OpenAI credentials configured in .env
- PostgreSQL with extracted_records table

Run with: pytest tests/knowledge/test_vector_store_integration.py -v -s
"""

import time
from uuid import uuid4

import pytest
from pymilvus import Collection, connections, utility

from src.core.config import settings
from src.extraction.models import ExtractedRecord
from src.knowledge.vector_store import VectorStore


@pytest.fixture(scope="module")
def milvus_connection():
    """Connect to Milvus for the test session."""
    connections.connect(
        alias="default",
        host=settings.milvus_host,
        port=settings.milvus_port,
    )
    yield
    connections.disconnect("default")


@pytest.fixture
def test_collection_name():
    """Use a test-specific collection name."""
    return f"test_dsol_records_{uuid4().hex[:8]}"


@pytest.fixture
async def vector_store_with_test_collection(test_collection_name):
    """Create VectorStore with test collection."""
    # Temporarily override collection name
    original_collection = settings.milvus_collection
    settings.milvus_collection = test_collection_name

    vector_store = VectorStore()
    await vector_store.create_collection()

    yield vector_store

    # Cleanup: drop test collection
    if utility.has_collection(test_collection_name):
        utility.drop_collection(test_collection_name)

    # Restore original collection name
    settings.milvus_collection = original_collection


@pytest.fixture
def sample_records():
    """Create sample ExtractedRecord objects for testing."""
    document_id = uuid4()

    records = [
        ExtractedRecord(
            content={
                "Item Code": "2024-4_0019",
                "Report": "not applicable",
                "Screen": "Applicable",
            },
            headers=["Item Code", "Report", "Screen"],
            _source={
                "document_id": str(document_id),
                "filename": "distribution_spec.xlsx",
                "sheet": "Item Mapping",
                "table_id": 1,
                "row": 10,
                "col_range": "A:D",
            },
        ),
        ExtractedRecord(
            content={
                "Item Code": "2024-4_0020",
                "Report": "applicable",
                "Screen": "Not Applicable",
            },
            headers=["Item Code", "Report", "Screen"],
            _source={
                "document_id": str(document_id),
                "filename": "distribution_spec.xlsx",
                "sheet": "Item Mapping",
                "table_id": 1,
                "row": 11,
                "col_range": "A:D",
            },
        ),
        ExtractedRecord(
            content={
                "Item Code": "2024-4_0021",
                "Report": "applicable",
                "Screen": "Applicable",
            },
            headers=["Item Code", "Report", "Screen"],
            _source={
                "document_id": str(document_id),
                "filename": "distribution_spec.xlsx",
                "sheet": "Item Mapping",
                "table_id": 1,
                "row": 12,
                "col_range": "A:D",
            },
        ),
    ]

    return document_id, records


@pytest.mark.integration
@pytest.mark.asyncio
async def test_create_collection(milvus_connection, test_collection_name):
    """Test creating Milvus collection with proper schema."""
    # Create vector store
    original_collection = settings.milvus_collection
    settings.milvus_collection = test_collection_name

    vector_store = VectorStore()
    await vector_store.create_collection()

    try:
        # Verify collection exists
        assert utility.has_collection(test_collection_name)

        # Get collection and verify schema
        collection = Collection(test_collection_name)
        schema = collection.schema

        # Verify fields
        field_names = [field.name for field in schema.fields]
        assert "id" in field_names
        assert "vector" in field_names
        assert "record_id" in field_names
        assert "document_id" in field_names
        assert "filename" in field_names
        assert "sheet_name" in field_names
        assert "row_number" in field_names
        assert "text_content" in field_names
        assert "has_symbols" in field_names

        # Verify vector dimensions
        vector_field = [f for f in schema.fields if f.name == "vector"][0]
        assert vector_field.params["dim"] == 3072

        # Verify index exists
        indexes = collection.indexes
        assert len(indexes) > 0
        vector_index = collection.index()
        assert vector_index.params["index_type"] == "HNSW"
        assert vector_index.params["metric_type"] == "COSINE"

    finally:
        # Cleanup
        if utility.has_collection(test_collection_name):
            utility.drop_collection(test_collection_name)
        settings.milvus_collection = original_collection


@pytest.mark.integration
@pytest.mark.asyncio
@pytest.mark.skipif(
    not settings.azure_openai_api_key,
    reason="Azure OpenAI credentials not configured",
)
async def test_index_records_end_to_end(
    vector_store_with_test_collection, sample_records, test_collection_name
):
    """Test end-to-end record indexing with real Azure OpenAI and Milvus."""
    vector_store = vector_store_with_test_collection
    document_id, records = sample_records

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

    # Index records
    start_time = time.time()
    count = await vector_store.index_records(records, document_id, record_ids)
    duration = time.time() - start_time

    # Verify indexing
    assert count == 3
    print(f"\nIndexed {count} records in {duration:.2f}s")

    # Wait for Milvus to persist data
    collection = Collection(test_collection_name)
    collection.load()

    # Verify vectors were stored
    result = collection.query(
        expr=f'document_id == "{str(document_id)}"',
        output_fields=["record_id", "filename", "text_content"],
    )

    assert len(result) == 3
    assert all(r["filename"] == "distribution_spec.xlsx" for r in result)


@pytest.mark.integration
@pytest.mark.asyncio
@pytest.mark.skipif(
    not settings.azure_openai_api_key,
    reason="Azure OpenAI credentials not configured",
)
async def test_index_large_batch_performance(
    vector_store_with_test_collection, test_collection_name
):
    """Test performance target: ≤ 30s per 1000 records."""
    vector_store = vector_store_with_test_collection
    document_id = uuid4()

    # Create 100 records for performance test (scaled down for CI)
    records = []
    record_ids = []

    for i in range(100):
        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": "performance_test.xlsx",
                "sheet": "Items",
                "table_id": 1,
                "row": i + 1,
                "col_range": "A:D",
            },
        )
        records.append(record)
        record_ids.append(uuid4())

    # Index records and measure time
    start_time = time.time()
    count = await vector_store.index_records(records, document_id, record_ids)
    duration = time.time() - start_time

    # Verify results
    assert count == 100
    records_per_second = count / duration

    print(f"\nPerformance metrics:")
    print(f"  - Records indexed: {count}")
    print(f"  - Duration: {duration:.2f}s")
    print(f"  - Records per second: {records_per_second:.2f}")
    print(f"  - Projected time for 1000 records: {(1000 / records_per_second):.2f}s")

    # Performance target: ≤ 30s per 1000 records
    # For 100 records, should be ≤ 3s
    assert duration <= 5.0, f"Indexing too slow: {duration:.2f}s for 100 records"


@pytest.mark.integration
@pytest.mark.asyncio
@pytest.mark.skipif(
    not settings.azure_openai_api_key,
    reason="Azure OpenAI credentials not configured",
)
async def test_similarity_search(
    vector_store_with_test_collection, sample_records, test_collection_name
):
    """Test similarity search retrieval after indexing."""
    vector_store = vector_store_with_test_collection
    document_id, records = sample_records

    # Index records
    record_ids = [uuid4() for _ in records]
    await vector_store.index_records(records, document_id, record_ids)

    # Wait for Milvus to load
    collection = Collection(test_collection_name)
    collection.load()

    # Perform similarity search
    from src.knowledge.embeddings import EmbeddingService

    embedding_service = EmbeddingService()

    # Search for "applicable report"
    query_text = "applicable report screening"
    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=["record_id", "text_content", "row_number"],
    )

    # Verify results
    assert len(search_results) > 0
    hits = search_results[0]
    assert len(hits) > 0

    print(f"\nSimilarity search results for '{query_text}':")
    for i, hit in enumerate(hits, 1):
        print(f"  {i}. Row {hit.entity.get('row_number')}: {hit.distance:.4f}")
        print(f"     {hit.entity.get('text_content')[:100]}...")

    # Verify results are relevant (contain "report" or "applicable")
    top_result = hits[0]
    text_content = top_result.entity.get("text_content", "")
    assert any(
        keyword in text_content.lower()
        for keyword in ["report", "applicable", "screen"]
    )


@pytest.mark.integration
@pytest.mark.asyncio
async def test_idempotent_indexing(
    vector_store_with_test_collection, sample_records, test_collection_name
):
    """Test that re-indexing same records is idempotent (upsert)."""
    vector_store = vector_store_with_test_collection
    document_id, records = sample_records

    record_ids = [uuid4() for _ in records]

    # Index records twice
    count1 = await vector_store.index_records(records, document_id, record_ids)
    count2 = await vector_store.index_records(records, document_id, record_ids)

    assert count1 == 3
    assert count2 == 3

    # Verify only 3 records exist (not 6)
    collection = Collection(test_collection_name)
    collection.load()

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

    # Should have exactly 3 records (upsert replaced duplicates)
    assert len(result) == 3


@pytest.mark.integration
@pytest.mark.asyncio
async def test_index_records_with_symbols(
    vector_store_with_test_collection, test_collection_name
):
    """Test indexing records with resolved symbol content."""
    vector_store = vector_store_with_test_collection
    document_id = uuid4()

    # Create record with symbol and resolved content
    record = ExtractedRecord(
        content={"Item": "TEST-001", "Status": "○"},
        headers=["Item", "Status"],
        _source={
            "document_id": str(document_id),
            "filename": "symbols_test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 1,
            "col_range": "A:B",
        },
    )

    # Add resolved content
    record.resolved_content = {
        "Item": "TEST-001",
        "Status": "○",
        "Status_resolved": "Applicable/Yes",
    }

    record_ids = [uuid4()]

    # Index record
    count = await vector_store.index_records([record], document_id, record_ids)
    assert count == 1

    # Verify has_symbols flag is set
    collection = Collection(test_collection_name)
    collection.load()

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

    assert len(result) == 1
    assert result[0]["has_symbols"] is True
    assert "Status_resolved: Applicable/Yes" in result[0]["text_content"]


@pytest.mark.integration
@pytest.mark.asyncio
async def test_collection_creation_idempotent(
    milvus_connection, test_collection_name
):
    """Test that creating collection twice doesn't fail."""
    original_collection = settings.milvus_collection
    settings.milvus_collection = test_collection_name

    try:
        vector_store = VectorStore()

        # Create collection twice
        await vector_store.create_collection()
        await vector_store.create_collection()  # Should not fail

        # Verify collection exists
        assert utility.has_collection(test_collection_name)

    finally:
        if utility.has_collection(test_collection_name):
            utility.drop_collection(test_collection_name)
        settings.milvus_collection = original_collection