"""Unit tests for StructuredStore."""

from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4

import pytest
from sqlalchemy.exc import IntegrityError, OperationalError

from src.core.exceptions import IndexingError
from src.extraction.models import ExtractedRecord
from src.knowledge.structured_store import StructuredStore


@pytest.mark.asyncio
async def test_store_records_success(db_session, sample_extracted_records, test_document_id):
    """Test storing records successfully in PostgreSQL."""
    # Arrange
    store = StructuredStore(db_session)

    # Mock _store_batch_with_retry to return count
    store._store_batch_with_retry = AsyncMock(return_value=10)

    # Act
    count = await store.store_records(sample_extracted_records, test_document_id)

    # Assert
    assert count == 10
    assert store._store_batch_with_retry.call_count == 1


@pytest.mark.asyncio
async def test_store_records_empty_list(db_session, test_document_id):
    """Test storing empty record list returns 0."""
    # Arrange
    store = StructuredStore(db_session)

    # Mock _store_batch_with_retry (should not be called for empty list)
    store._store_batch_with_retry = AsyncMock()

    # Act
    count = await store.store_records([], test_document_id)

    # Assert
    assert count == 0
    assert store._store_batch_with_retry.call_count == 0


@pytest.mark.asyncio
async def test_store_records_batch_splitting():
    """Test that records are split into batches of 500."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)

    # Create 1200 mock records (should be 3 batches)
    test_doc_id = uuid4()
    records = []
    for i in range(1200):
        record = ExtractedRecord(
            content={"Item": f"Item {i}"},
            headers=["Item"],
            _source={
                "document_id": str(test_doc_id),
                "filename": "test.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": i + 1,
                "col_range": "A:A",
            },
        )
        records.append(record)

    # Mock the _store_batch_with_retry method
    store._store_batch_with_retry = AsyncMock(return_value=500)

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

    # Assert
    # Should be called 3 times (1200 / 500 = 2.4 → 3 batches)
    assert store._store_batch_with_retry.call_count == 3
    # First two batches: 500 records each, last batch: 200 records
    assert store._store_batch_with_retry.call_args_list[0][0][0].__len__() == 500
    assert store._store_batch_with_retry.call_args_list[1][0][0].__len__() == 500
    assert store._store_batch_with_retry.call_args_list[2][0][0].__len__() == 200


@pytest.mark.asyncio
async def test_store_records_raises_indexing_error_on_failure():
    """Test that storage failures raise IndexingError."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)
    test_doc_id = uuid4()

    record = ExtractedRecord(
        content={"Item": "Test"},
        headers=["Item"],
        _source={
            "document_id": str(test_doc_id),
            "filename": "test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 1,
            "col_range": "A:A",
        },
    )

    # Mock _store_batch_with_retry to raise exception
    store._store_batch_with_retry = AsyncMock(
        side_effect=OperationalError("Database connection lost", None, None)
    )

    # Act & Assert
    with pytest.raises(IndexingError) as exc_info:
        await store.store_records([record], test_doc_id)

    assert "Failed to store records" in str(exc_info.value)
    assert exc_info.value.code == "INDEXING_ERROR"


@pytest.mark.asyncio
async def test_store_batch_with_retry_success_first_attempt():
    """Test successful storage on first attempt (no retries needed)."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)
    test_doc_id = uuid4()

    record = ExtractedRecord(
        content={"Item": "Test"},
        headers=["Item"],
        _source={
            "document_id": str(test_doc_id),
            "filename": "test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 1,
            "col_range": "A:A",
        },
    )

    # Mock _store_batch to succeed
    store._store_batch = AsyncMock(return_value=1)

    # Act
    count = await store._store_batch_with_retry([record], test_doc_id, batch_idx=0)

    # Assert
    assert count == 1
    assert store._store_batch.call_count == 1


@pytest.mark.asyncio
async def test_store_batch_with_retry_retries_on_operational_error():
    """Test retry logic on transient OperationalError."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)
    test_doc_id = uuid4()

    record = ExtractedRecord(
        content={"Item": "Test"},
        headers=["Item"],
        _source={
            "document_id": str(test_doc_id),
            "filename": "test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 1,
            "col_range": "A:A",
        },
    )

    # Mock _store_batch to fail twice, then succeed
    store._store_batch = AsyncMock(
        side_effect=[
            OperationalError("Connection lost", None, None),
            OperationalError("Connection lost", None, None),
            1,  # Success on third attempt
        ]
    )

    # Mock asyncio.sleep to avoid actual delays
    with patch("asyncio.sleep", new_callable=AsyncMock):
        # Act
        count = await store._store_batch_with_retry([record], test_doc_id, batch_idx=0)

    # Assert
    assert count == 1
    assert store._store_batch.call_count == 3


@pytest.mark.asyncio
async def test_store_batch_with_retry_fails_after_max_retries():
    """Test that retry logic gives up after MAX_RETRIES attempts."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)
    test_doc_id = uuid4()

    record = ExtractedRecord(
        content={"Item": "Test"},
        headers=["Item"],
        _source={
            "document_id": str(test_doc_id),
            "filename": "test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 1,
            "col_range": "A:A",
        },
    )

    # Mock _store_batch to always fail
    store._store_batch = AsyncMock(
        side_effect=OperationalError("Persistent error", None, None)
    )

    # Mock asyncio.sleep to avoid actual delays
    with patch("asyncio.sleep", new_callable=AsyncMock):
        # Act & Assert
        with pytest.raises(OperationalError):
            await store._store_batch_with_retry([record], test_doc_id, batch_idx=0)

    # Should try 3 times (MAX_RETRIES)
    assert store._store_batch.call_count == 3


@pytest.mark.asyncio
async def test_store_batch_with_retry_no_retry_on_integrity_error():
    """Test that IntegrityError is not retried (permanent error)."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)
    test_doc_id = uuid4()

    record = ExtractedRecord(
        content={"Item": "Test"},
        headers=["Item"],
        _source={
            "document_id": str(test_doc_id),
            "filename": "test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 1,
            "col_range": "A:A",
        },
    )

    # Mock _store_batch to raise IntegrityError
    store._store_batch = AsyncMock(
        side_effect=IntegrityError("Duplicate key", None, None)
    )

    # Act & Assert
    with pytest.raises(IntegrityError):
        await store._store_batch_with_retry([record], test_doc_id, batch_idx=0)

    # Should only try once (no retries for IntegrityError)
    assert store._store_batch.call_count == 1


@pytest.mark.asyncio
async def test_to_db_model_converts_correctly(test_document_id):
    """Test conversion of ExtractedRecord to database model dict."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)

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

    # Act
    db_record = store._to_db_model(record, test_document_id)

    # Assert
    assert db_record["document_id"] == test_document_id
    assert db_record["sheet_name"] == "Sheet1"
    assert db_record["row_number"] == 10
    assert db_record["col_range"] == "A:C"
    assert db_record["content"] == record.content
    assert db_record["headers"] == record.headers
    assert "resolved_content" not in db_record  # Not present in this record


@pytest.mark.asyncio
async def test_to_db_model_includes_resolved_content(test_document_id):
    """Test that resolved_content is included if present."""
    # Arrange
    mock_session = MagicMock()
    store = StructuredStore(mock_session)

    record = ExtractedRecord(
        content={
            "Item Code": "2024-4_0001",
            "Status": "○",
        },
        headers=["Item Code", "Status"],
        _source={
            "document_id": str(test_document_id),
            "filename": "test.xlsx",
            "sheet": "Sheet1",
            "table_id": 1,
            "row": 10,
            "col_range": "A:B",
        },
    )

    # Add resolved_content (simulating symbol resolution)
    record.resolved_content = {
        "Item Code": "2024-4_0001",
        "Status": "Available",
    }

    # Act
    db_record = store._to_db_model(record, test_document_id)

    # Assert
    assert db_record["resolved_content"] == record.resolved_content