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

from __future__ import annotations

import uuid
from typing import Any, Iterable, List

import pandas as pd
from azure.cosmos import CosmosClient


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")

        client = CosmosClient(endpoint, credential=key)
        database = client.get_database_client(database_name)
        self._container = database.get_container_client(container_name)

    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")

        inserted = 0
        for _, row in df.iterrows():
            document = {
                "id": str(uuid.uuid4()),
                partition_key_column: row[partition_key_column],
                **row.to_dict(),
            }
            self._container.create_item(body=document)
            inserted += 1
        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:
            items_iter: Iterable[dict[str, Any]]
            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,
                )
            items = list(items_iter)
        else:
            if partition_key_value is None:
                raise ValueError("partition_key_value is required when partition_key_column is provided")
            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,
                )
            )

        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
