"""Unit tests for the CosmosManager helper."""

import logging

import pandas as pd
import pytest

import hel_elmuro.cosmos_manager as cosmos_module
from hel_elmuro.cosmos_manager import CosmosManager


# Logging configuration is provided by pytest.ini.
logger = logging.getLogger(__name__)


class _StubContainer:
    def __init__(self) -> None:
        self.items = []

    def create_item(self, body):  # type: ignore[no-untyped-def]
        self.items.append(body)

    def read_all_items(self):  # type: ignore[no-untyped-def]
        return list(self.items)

    def query_items(self, query, parameters=None, enable_cross_partition_query=False):  # type: ignore[no-untyped-def]
        if not parameters:
            return list(self.items)

        column = query.split("c.")[1].split()[0]
        value = parameters[0]["value"]
        return [item for item in self.items if item.get(column) == value]


class _StubDatabase:
    def __init__(self, container: _StubContainer) -> None:
        self.container = container
        self.requested_container_name = None

    def get_container_client(self, name):  # type: ignore[no-untyped-def]
        self.requested_container_name = name
        return self.container


class _StubCosmosClient:
    def __init__(self, endpoint: str, credential: str) -> None:
        self.endpoint = endpoint
        self.credential = credential
        self.requested_database = None
        self._container = _StubContainer()

    def get_database_client(self, database_name):  # type: ignore[no-untyped-def]
        self.requested_database = database_name
        return _StubDatabase(self._container)


@pytest.fixture
def cosmos_manager(monkeypatch):
    """Return a CosmosManager wired to a stubbed Cosmos client."""
    logger.info("Creating stub Cosmos client and manager")
    stub_client = _StubCosmosClient("https://test", "key")
    monkeypatch.setattr(cosmos_module, "CosmosClient", lambda *a, **kw: stub_client)
    manager = CosmosManager("https://test", "key", "elmuro", "drivers")
    return manager, stub_client


def test_save_dataframe_requires_partition_key(cosmos_manager):
    """Verify that the helper raises an error when the partition key column is missing."""
    manager, _stub_client = cosmos_manager

    # 1) Provide data missing the required partition key column.
    df = pd.DataFrame([{"name": "Anne", "hours": 6}])

    with pytest.raises(KeyError):
        # 2) Verify the helper raises the expected exception before writing anything.
        logger.info("Invoking save_dataframe expecting KeyError")
        manager.save_dataframe(df, "driver_id")



def test_fetch_documents_returns_all_items(cosmos_manager):
    """Verify saving two rows succeeds, fetch_documents returns them, and filtering by partition key works."""
    manager, stub_client = cosmos_manager

    # 1) Seed the container with known data via save_dataframe and assert the write path works.
    df = pd.DataFrame(
        [
            {"driver_id": "D1", "name": "Anne", "hours": 6},
            {"driver_id": "D2", "name": "Ben", "hours": 7},
        ]
    )

    logger.info("Invoking save_dataframe")
    inserted = manager.save_dataframe(df, "driver_id")

    logger.info("Inserted %s documents", inserted)
    assert inserted == 2
    assert len(stub_client._container.items) == 2
    assert {item["driver_id"] for item in stub_client._container.items} == {"D1", "D2"}

    # 2) Retrieve all persisted documents using the new API and assert the read path matches the written data.
    logger.info("Invoking fetch_documents")
    retrieved = manager.fetch_documents()

    logger.info("Fetched %s documents", len(retrieved))
    assert len(retrieved) == 2
    assert {item["driver_id"] for item in retrieved} == {"D1", "D2"}

    # 3) Fetch documents filtered by one partition key value and validate the outcome.
    logger.info("Invoking fetch_documents with filter")
    filtered = manager.fetch_documents(
        partition_key_column="driver_id", partition_key_value="D2"
    )

    assert len(filtered) == 1
    assert filtered[0]["driver_id"] == "D2"

