"""Cosmos DB helper utilities for ElMuro workflows."""

from __future__ import annotations

import logging
import uuid
from typing import Any, Iterable, List

import pandas as pd
from azure.cosmos import CosmosClient
from azure.cosmos.exceptions import CosmosHttpResponseError


logger = logging.getLogger(__name__)
for _logger_name in ("azure", "azure.core", "azure.cosmos"):
    sdk_logger = logging.getLogger(_logger_name)
    if sdk_logger.level == logging.NOTSET:
        sdk_logger.setLevel(logging.WARNING)


class CosmosManager:
    """Persist pandas data into Cosmos DB containers."""

    def __init__(self, endpoint: str, key: str, database_name: str, container_name: str) -> None:
        if not endpoint:
            raise ValueError("endpoint cannot be empty")
        if not key:
            raise ValueError("key cannot be empty")
        if not database_name:
            raise ValueError("database_name cannot be empty")
        if not container_name:
            raise ValueError("container_name cannot be empty")

        logger.info("Connecting to Cosmos DB endpoint %s (database=%s, container=%s).", endpoint, database_name, container_name)
        try:
            client = CosmosClient(endpoint, credential=key)
            database = client.get_database_client(database_name)
            self._container = database.get_container_client(container_name)
        except CosmosHttpResponseError as exc:
            logger.error("Cosmos DB connection failed: %s", exc)
            raise RuntimeError("Failed to connect to Cosmos DB.") from exc
        except Exception as exc:  # noqa: BLE001 - surface unexpected issues
            logger.error("Unexpected error while connecting to Cosmos DB: %s", exc)
            raise RuntimeError("Unexpected error while connecting to Cosmos DB.") from exc
        self._container_name = container_name
        logger.info("Cosmos container client ready.")

    def save_dataframe(self, df: pd.DataFrame, partition_key_column: str) -> int:
        if partition_key_column not in df.columns:
            raise KeyError(f"partition key column '{partition_key_column}' not found")

        total_rows = len(df)
        logger.info("Persisting %d records into Cosmos container %s.", total_rows, self._container_name)
        inserted = 0
        for index, row in df.iterrows():
            document = {
                "id": str(uuid.uuid4()),
                partition_key_column: row[partition_key_column],
                **row.to_dict(),
            }
            try:
                self._container.create_item(body=document)
            except CosmosHttpResponseError as exc:
                logger.error("Failed to create document in Cosmos DB: %s", exc)
                raise RuntimeError("Failed to save document to Cosmos DB.") from exc
            inserted += 1
            if (index + 1) % 100 == 0 or index + 1 == total_rows:
                logger.info("Inserted %d/%d records into Cosmos container.", index + 1, total_rows)
        logger.info("Inserted %d records into Cosmos container.", inserted)
        return inserted

    def fetch_documents(
        self,
        partition_key_column: str | None = None,
        partition_key_value: Any | None = None,
    ) -> List[dict[str, Any]]:
        """Return documents from the configured container, optionally filtered by partition key."""

        if partition_key_column is None:
            logger.info("Fetching all documents from Cosmos container.")
            items_iter: Iterable[dict[str, Any]]
            try:
                if hasattr(self._container, "read_all_items"):
                    items_iter = self._container.read_all_items()
                else:
                    items_iter = self._container.query_items(
                        query="SELECT * FROM c",
                        parameters=[],
                        enable_cross_partition_query=True,
                    )
            except CosmosHttpResponseError as exc:
                logger.error("Failed to read documents from Cosmos DB: %s", exc)
                raise RuntimeError("Failed to read documents from Cosmos DB.") from exc
            items = list(items_iter)
        else:
            if partition_key_value is None:
                raise ValueError("partition_key_value is required when partition_key_column is provided")
            logger.info(
                "Querying Cosmos container by partition %s=%s.",
                partition_key_column,
                partition_key_value,
            )
            try:
                items = list(
                    self._container.query_items(
                        query=f"SELECT * FROM c WHERE c.{partition_key_column} = @pk",
                        parameters=[{"name": "@pk", "value": partition_key_value}],
                        enable_cross_partition_query=True,
                    )
                )
            except CosmosHttpResponseError as exc:
                logger.error("Failed to query documents from Cosmos DB: %s", exc)
                raise RuntimeError("Failed to query documents from Cosmos DB.") from exc

        logger.info("Retrieved %d documents from Cosmos container %s.", len(items), self._container_name)
        if partition_key_column and partition_key_value is not None:
            return [item for item in items if item.get(partition_key_column) == partition_key_value]
        return items
