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

from __future__ import annotations

import logging
from typing import Iterable, Tuple

import pandas as pd
import pytest

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


logger = logging.getLogger(__name__)


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

    def create_item(self, body: dict, **kwargs) -> dict:
        self.items.append(body)
        return body

    def read_all_items(self) -> Iterable[dict]:
        return list(self.items)

    def query_items(self, query: str, parameters: Iterable[dict], **kwargs) -> Iterable[dict]:
        partition_key = parameters[0]["value"]
        return [item for item in self.items if item.get("driver_id") == partition_key]


class StubDatabase:
    def __init__(self) -> None:
        self._container = StubContainer()

    def get_container_client(self, container_name: str) -> StubContainer:  # noqa: ARG002
        return self._container


class StubCosmosClient:
    def __init__(self) -> None:
        self._database = StubDatabase()

    def get_database_client(self, database_name: str) -> StubDatabase:  # noqa: ARG002
        return self._database


@pytest.fixture()
def cosmos_manager(monkeypatch: pytest.MonkeyPatch) -> Tuple[CosmosManager, StubCosmosClient]:
    stub_client = StubCosmosClient()
    monkeypatch.setattr(cosmos_module, "CosmosClient", lambda *args, **kwargs: stub_client)
    manager = CosmosManager("https://test", "key", "elmuro", "drivers")
    return manager, stub_client


def test_cosmos_manager_workflow(cosmos_manager: Tuple[CosmosManager, StubCosmosClient]) -> None:
    """End-to-end save/fetch workflow including validation of required partition keys."""
    manager, stub_client = cosmos_manager

    logger.info("Scenario: verify missing partition key rejection and full fetch cycle")

    df_missing_key = pd.DataFrame([{"name": "Anne", "hours": 6}])
    with pytest.raises(KeyError):
        logger.info(
            "Attempting save_dataframe with %s rows; partition key column 'driver_id' is absent",
            len(df_missing_key),
        )
        manager.save_dataframe(df_missing_key, "driver_id")

    df = pd.DataFrame(
        [
            {"driver_id": "D1", "name": "Anne", "hours": 6},
            {"driver_id": "D2", "name": "Ben", "hours": 7},
        ]
    )

    logger.info(
        "Attempting save_dataframe with %s rows using partition key 'driver_id'",
        len(df),
    )
    inserted = manager.save_dataframe(df, "driver_id")
    assert inserted == 2
    assert len(stub_client._database._container.items) == 2
    assert {item["driver_id"] for item in stub_client._database._container.items} == {"D1", "D2"}

    logger.info(
        "Attempting fetch_documents filtered by partition key 'driver_id' == '%s' (covers unfiltered path via stub setup)",
        "D2",
    )
    filtered = manager.fetch_documents(
        partition_key_column="driver_id", partition_key_value="D2"
    )

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

