"""Integration tests for StructuredStore with real database.

These tests require a PostgreSQL database connection.
Set TEST_DATABASE_URL environment variable to run against a test database.
"""

import time
from uuid import uuid4

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

from src.db.models import Base, Document, ExtractedRecord as ExtractedRecordModel
from src.extraction.models import ExtractedRecord
from src.knowledge.structured_store import StructuredStore

# Mark all tests in this module as integration tests
pytestmark = pytest.mark.integration


@pytest.fixture(scope="module")
async def integration_db_engine():
    """Create a test database engine for integration tests.

    Uses PostgreSQL test database (or skip if not available).
    """
    import os

    test_db_url = os.getenv("TEST_DATABASE_URL")
    if not test_db_url:
        pytest.skip("TEST_DATABASE_URL not set - skipping integration tests")

    engine = create_async_engine(
        test_db_url,
        echo=False,
        pool_pre_ping=True,
    )

    # Create all tables
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)
        await conn.run_sync(Base.metadata.create_all)

    yield engine

    # Cleanup
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)

    await engine.dispose()


@pytest.fixture
async def integration_db_session(integration_db_engine):
    """Create a test database session for integration tests."""
    async_session_maker = sessionmaker(
        integration_db_engine,
        class_=AsyncSession,
        expire_on_commit=False,
    )

    async with async_session_maker() as session:
        yield session
        await session.rollback()
        await session.close()


@pytest.fixture
async def test_document(integration_db_session):
    """Create a test document in the database."""
    document = Document(
        id=uuid4(),
        filename="test_integration.xlsx",
        file_path="/tmp/test_integration.xlsx",
        status="pending",
    )

    integration_db_session.add(document)
    await integration_db_session.commit()
    await integration_db_session.refresh(document)

    return document


@pytest.mark.asyncio
async def test_store_records_integration(integration_db_session, test_document):
    """Test storing records in real PostgreSQL database."""
    # Arrange
    store = StructuredStore(integration_db_session)

    records = []
    for i in range(10):
        record = ExtractedRecord(
            content={
                "Item Code": f"2024-4_{i:04d}",
                "Description": f"Test Item {i}",
                "Quantity": str(100 + i),
                "Screen": "○" if i % 2 == 0 else "×",
            },
            headers=["Item Code", "Description", "Quantity", "Screen"],
            _source={
                "document_id": str(test_document.id),
                "filename": test_document.filename,
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 10 + i,
                "col_range": "A:D",
            },
        )
        records.append(record)

    # Act
    count = await store.store_records(records, test_document.id)

    # Assert
    assert count == 10

    # Verify records in database
    result = await integration_db_session.execute(
        select(ExtractedRecordModel).where(
            ExtractedRecordModel.document_id == test_document.id
        )
    )
    stored_records = result.scalars().all()

    assert len(stored_records) == 10

    # Verify first record content
    first_record = stored_records[0]
    assert first_record.sheet_name == "Sheet1"
    assert first_record.row_number >= 10
    assert first_record.col_range == "A:D"
    assert "Item Code" in first_record.content
    assert first_record.headers == ["Item Code", "Description", "Quantity", "Screen"]


@pytest.mark.asyncio
async def test_store_records_duplicate_handling(integration_db_session, test_document):
    """Test ON CONFLICT handling for duplicate records."""
    # Arrange
    store = StructuredStore(integration_db_session)

    record = ExtractedRecord(
        content={
            "Item Code": "2024-4_0001",
            "Description": "Original Description",
            "Quantity": "100",
        },
        headers=["Item Code", "Description", "Quantity"],
        _source={
            "document_id": str(test_document.id),
            "filename": test_document.filename,
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 10,
            "col_range": "A:C",
        },
    )

    # Act - First insert
    count1 = await store.store_records([record], test_document.id)
    assert count1 == 1

    # Modify record content
    record.content["Description"] = "Updated Description"
    record.content["Quantity"] = "200"

    # Act - Second insert (should update, not fail)
    count2 = await store.store_records([record], test_document.id)
    assert count2 == 1

    # Assert - Verify only one record exists with updated content
    result = await integration_db_session.execute(
        select(ExtractedRecordModel).where(
            ExtractedRecordModel.document_id == test_document.id,
            ExtractedRecordModel.sheet_name == "Sheet1",
            ExtractedRecordModel.row_number == 10,
        )
    )
    stored_records = result.scalars().all()

    assert len(stored_records) == 1
    assert stored_records[0].content["Description"] == "Updated Description"
    assert stored_records[0].content["Quantity"] == "200"


@pytest.mark.asyncio
async def test_store_1000_records_performance(integration_db_session, test_document):
    """Test performance benchmark: 1000 records in < 3 seconds."""
    # Arrange
    store = StructuredStore(integration_db_session)

    # Create 1000 mock records
    records = []
    for i in range(1000):
        record = ExtractedRecord(
            content={
                "Item Code": f"2024-4_{i:04d}",
                "Description": f"Benchmark Item {i}",
                "Quantity": str(100 + i),
                "Screen": "○",
            },
            headers=["Item Code", "Description", "Quantity", "Screen"],
            _source={
                "document_id": str(test_document.id),
                "filename": test_document.filename,
                "sheet": "BenchmarkSheet",
                "table_id": 1,
                "row": 10 + i,
                "col_range": "A:D",
            },
        )
        records.append(record)

    # Act
    start_time = time.time()
    count = await store.store_records(records, test_document.id)
    duration = time.time() - start_time

    # Assert
    assert count == 1000
    assert duration < 3.0, f"Performance target missed: {duration:.2f}s (target: <3s)"

    # Verify all records stored
    result = await integration_db_session.execute(
        select(ExtractedRecordModel).where(
            ExtractedRecordModel.document_id == test_document.id,
            ExtractedRecordModel.sheet_name == "BenchmarkSheet",
        )
    )
    stored_records = result.scalars().all()
    assert len(stored_records) == 1000


@pytest.mark.asyncio
async def test_query_performance_exact_lookup(integration_db_session, test_document):
    """Test exact item code lookup performance: < 50ms."""
    # Arrange
    store = StructuredStore(integration_db_session)

    # Create 1000 records
    records = []
    for i in range(1000):
        record = ExtractedRecord(
            content={
                "Item Code": f"2024-4_{i:04d}",
                "Description": f"Query Test Item {i}",
            },
            headers=["Item Code", "Description"],
            _source={
                "document_id": str(test_document.id),
                "filename": test_document.filename,
                "sheet": "QuerySheet",
                "table_id": 1,
                "row": 10 + i,
                "col_range": "A:B",
            },
        )
        records.append(record)

    await store.store_records(records, test_document.id)

    # Act - Query by exact item code
    start_time = time.time()
    result = await integration_db_session.execute(
        select(ExtractedRecordModel).where(
            ExtractedRecordModel.document_id == test_document.id,
            ExtractedRecordModel.content["Item Code"].astext == "2024-4_0500",
        )
    )
    query_duration = (time.time() - start_time) * 1000  # Convert to ms

    # Assert
    found_records = result.scalars().all()
    assert len(found_records) == 1
    assert found_records[0].content["Item Code"] == "2024-4_0500"
    assert query_duration < 50, f"Query too slow: {query_duration:.2f}ms (target: <50ms)"


@pytest.mark.asyncio
async def test_query_performance_jsonb_containment(integration_db_session, test_document):
    """Test JSONB containment query performance: < 100ms."""
    # Arrange
    store = StructuredStore(integration_db_session)

    # Create 1000 records with Screen field
    records = []
    for i in range(1000):
        record = ExtractedRecord(
            content={
                "Item Code": f"2024-4_{i:04d}",
                "Screen": "○" if i % 3 == 0 else "×",
            },
            headers=["Item Code", "Screen"],
            _source={
                "document_id": str(test_document.id),
                "filename": test_document.filename,
                "sheet": "ContainmentSheet",
                "table_id": 1,
                "row": 10 + i,
                "col_range": "A:B",
            },
        )
        records.append(record)

    await store.store_records(records, test_document.id)

    # Act - JSONB containment query
    start_time = time.time()
    result = await integration_db_session.execute(
        select(ExtractedRecordModel).where(
            ExtractedRecordModel.document_id == test_document.id,
            ExtractedRecordModel.content.contains({"Screen": "○"}),
        )
    )
    query_duration = (time.time() - start_time) * 1000  # Convert to ms

    # Assert
    found_records = result.scalars().all()
    expected_count = len([r for r in records if r.content["Screen"] == "○"])
    assert len(found_records) == expected_count
    assert query_duration < 100, f"Query too slow: {query_duration:.2f}ms (target: <100ms)"


@pytest.mark.asyncio
async def test_query_performance_sheet_filter(integration_db_session, test_document):
    """Test sheet filter query performance: < 20ms."""
    # Arrange
    store = StructuredStore(integration_db_session)

    # Create 100 records in specific sheet
    records = []
    for i in range(100):
        record = ExtractedRecord(
            content={"Item": f"Item {i}"},
            headers=["Item"],
            _source={
                "document_id": str(test_document.id),
                "filename": test_document.filename,
                "sheet": "FastSheet",
                "table_id": 1,
                "row": 10 + i,
                "col_range": "A:A",
            },
        )
        records.append(record)

    await store.store_records(records, test_document.id)

    # Act - Query by document_id and sheet_name
    start_time = time.time()
    result = await integration_db_session.execute(
        select(ExtractedRecordModel).where(
            ExtractedRecordModel.document_id == test_document.id,
            ExtractedRecordModel.sheet_name == "FastSheet",
        )
    )
    query_duration = (time.time() - start_time) * 1000  # Convert to ms

    # Assert
    found_records = result.scalars().all()
    assert len(found_records) == 100
    assert query_duration < 20, f"Query too slow: {query_duration:.2f}ms (target: <20ms)"
