"""MongoDB database service with CRUD operations and index management.

This module provides TinyDBService, a high-level MongoDB client for the TinyAI
microservices platform. It handles common database operations with automatic
index management and ObjectId conversion.

Example:
  Basic usage with environment variables::

    from tinysdk import TinyDBService

    db = TinyDBService()
    doc_id = db.create("users", {"name": "John", "email": "john@example.com"})
"""

import os
from typing import Any, Dict, List, Optional, Tuple

from bson.objectid import ObjectId
from pymongo import ASCENDING, MongoClient
from pymongo.errors import AutoReconnect, NetworkTimeout, NotPrimaryError, ServerSelectionTimeoutError

from tinysdk.tinylogger import TinyLogger
from tinysdk.utils.decorators import exception_logger, retry

# Default index configuration for common collections
DEFAULT_INDEX_CONFIG: Dict[str, List[Any]] = {
  "users": [[("email", ASCENDING), {"unique": True}]],
  "documents": [[("fileName", ASCENDING), {"unique": True}]],
}

_db_retry_logger = TinyLogger()
_mongo_retry = retry(
  exceptions=(AutoReconnect, NetworkTimeout, NotPrimaryError, ServerSelectionTimeoutError),
  tries=3,
  delay=1,
  backoff=2,
  jitter=(0, 1),
  logger=_db_retry_logger,
)


class TinyDBService:
  """MongoDB service for CRUD operations with automatic index management.

  TinyDBService provides a high-level interface for MongoDB operations including
  create, read, update, delete, and bulk operations. It automatically manages
  indexes based on configuration and handles ObjectId conversions.

  Attributes:
    db_link (str): MongoDB connection string
    db_name (str): Database name
    client (MongoClient): Active MongoDB client connection
    db: MongoDB database instance
    index_config (Dict[str, List]): Index configuration for collections

  Example:
    >>> # Using environment variables
    >>> db = TinyDBService()
    >>>
    >>> # Or explicit parameters
    >>> db = TinyDBService(
    ...     db_link="mongodb://localhost:27017",
    ...     db_name="mydb"
    ... )
    >>>
    >>> # Create a document
    >>> doc_id = db.create("users", {"name": "John", "email": "john@example.com"})
  """

  @exception_logger
  def __init__(
    self,
    db_link: Optional[str] = None,
    db_name: Optional[str] = None,
    index_config: Optional[Dict[str, List[Any]]] = None,
    logger: Optional[TinyLogger] = None,
    connection_timeout_ms: int = 5000,
  ) -> None:
    """Initialize TinyDBService with MongoDB connection.

    Args:
      db_link (str, optional): MongoDB connection string. Defaults to DB_LINK env var.
      db_name (str, optional): Database name. Defaults to DB_NAME env var.
      index_config (Dict[str, List], optional): Custom index configuration to merge with defaults.

    Raises:
      AssertionError: If db_link or db_name not provided via parameters or environment variables.

    Environment Variables:
      DB_LINK: MongoDB connection string (e.g., mongodb://mongo-server:27017)
      DB_NAME: Database name (e.g., tinyDatabase)
    """
    db_link = db_link or os.getenv("DB_LINK")
    db_name = db_name or os.getenv("DB_NAME")
    assert db_link, "DB_LINK must be provided via parameter or environment variable"
    assert db_name, "DB_NAME must be provided via parameter or environment variable"
    self.db_link: str = db_link
    self.db_name: str = db_name

    self.logger: TinyLogger = logger or TinyLogger()
    self.client: MongoClient = MongoClient(self.db_link, serverSelectionTimeoutMS=connection_timeout_ms)
    self.db = self.client.get_database(self.db_name)

    # Initialize index configuration
    self.index_config = DEFAULT_INDEX_CONFIG.copy()
    self._initialized_collections: set = set()
    if index_config:
      self.index_config.update(index_config)

  @_mongo_retry
  @exception_logger
  def create(self, collection: str, document: Dict[str, Any]) -> str:
    """
    Inserts a new document into the specified collection.

    Args:
        collection (str): The name of the collection to insert the document into.
        document (Dict[str, Any]): The document to insert.

    Returns:
        str: The inserted document ID as a string.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    if document.get("_id") is not None and not isinstance(document.get("_id"), str):
      pass  # already ObjectId or other valid type — leave as-is
    elif isinstance(document.get("_id"), str):
      document = {**document, "_id": ObjectId(document["_id"])}
    response = col.insert_one(document)
    return str(response.inserted_id)

  @_mongo_retry
  @exception_logger
  def read(
    self,
    collection: str,
    filter: Optional[Dict[str, Any]] = None,
    sort: Optional[Tuple[str, int]] = None,
    page: int = 0,
    limit: int = 0,
  ) -> List[Dict[str, Any]]:
    """
    Retrieves documents from the specified collection based on filters and options.

    Args:
        collection (str): The name of the collection to query.
        filter (Dict[str, Any], optional): Filter criteria. Defaults to no filter.
        sort (Tuple[str, int], optional): A (field, direction) tuple for sorting.
        page (int): Page number for pagination (used with limit).
        limit (int): Maximum number of documents to return (0 means no limit).

    Returns:
        List[Dict[str, Any]]: List of documents with _id converted to string.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    search_filter = self._convert_filter_id(filter or {})
    cursor = col.find(search_filter)
    if sort is not None:
      cursor = cursor.sort(sort[0], sort[1])
    if limit > 0:
      cursor = cursor.skip(page * limit).limit(limit)
    return [{**doc, "_id": str(doc.get("_id"))} for doc in cursor]

  @_mongo_retry
  @exception_logger
  def update(self, collection: str, filter: Dict[str, Any], data: Dict[str, Any]) -> int:
    """
    Updates documents in the specified collection based on the filter.

    Args:
        collection (str): The name of the collection to update.
        filter (Dict[str, Any]): Filter criteria to identify documents to update.
        data (Dict[str, Any]): Update operations to apply (e.g. {"$set": {...}}).

    Returns:
        int: The number of modified documents.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    search_filter = self._convert_filter_id(filter)
    response = col.update_many(search_filter, data)
    return response.modified_count

  @_mongo_retry
  @exception_logger
  def delete(self, collection: str, filter: Optional[Dict[str, Any]] = None) -> int:
    """
    Deletes documents from the specified collection based on the filter.

    Args:
        collection (str): The name of the collection to delete from.
        filter (Dict[str, Any], optional): Filter criteria. Defaults to {} (delete all).

    Returns:
        int: The number of deleted documents.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    search_filter = self._convert_filter_id(filter or {})
    response = col.delete_many(search_filter)
    return response.deleted_count

  @exception_logger
  def _init_index(self, collection_name: str) -> None:
    """
    Initializes indexes for the specified collection based on the index configuration.

    Args:
        collection_name (str): The name of the collection for which to create indexes.
    """
    assert collection_name is not None, "Collection name cannot be None"
    if collection_name in self._initialized_collections:
      return
    collection = self.db.get_collection(collection_name)
    # Iterate through the indexes for this collection
    for index in self.index_config.get(collection_name, []):
      # The index specification (a list) can have additional parameters
      fields = index[:-1] if isinstance(index[-1], dict) else index
      options = index[-1] if isinstance(index[-1], dict) else {}
      # Create the index
      collection.create_index(fields, **options)
    self._initialized_collections.add(collection_name)

  @exception_logger
  def _convert_filter_id(self, filter: Dict[str, Any]) -> Dict[str, Any]:
    """
    Converts string IDs in the filter to ObjectId.

    Args:
    filter (Dict[str, Any]): Dictionary containing filter criteria.

    Returns:
    Dict[str, Any]: Filter with ObjectId values.
    """
    filter = filter.copy()  # do not mutate caller's dict
    id = filter.get("_id", None)
    if isinstance(id, str):
      filter["_id"] = ObjectId(id)
    elif isinstance(id, dict) and "$in" in id:
      filter["_id"]["$in"] = [ObjectId(id_str) for id_str in id["$in"]]
    return filter

  @_mongo_retry
  @exception_logger
  def create_multiple(self, collection: str, documents: List[Dict[str, Any]]) -> List[str]:
    """
    Inserts multiple documents into the specified collection.

    Args:
        collection (str): The name of the collection to insert documents into.
        documents (List[Dict[str, Any]]): The list of documents to insert.

    Returns:
        List[str]: The inserted document IDs as strings.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    result = col.insert_many(documents)
    return [str(id) for id in result.inserted_ids]

  @_mongo_retry
  @exception_logger
  def exists(self, collection: str, filter: Optional[Dict[str, Any]] = None) -> bool:
    """
    Checks if a document exists in the specified collection based on the filter.

    Args:
        collection (str): The name of the collection to check.
        filter (Dict[str, Any], optional): Filter criteria. Defaults to {} (any document).

    Returns:
        bool: True if a matching document exists, False otherwise.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    search_filter = self._convert_filter_id(filter or {})
    return col.count_documents(search_filter) > 0

  @exception_logger
  def close(self) -> None:
    """Closes the MongoDB client connection and releases the connection pool."""
    self.client.close()
    self.logger.log("MongoDB client connection closed.")
