"""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 MongoClient
from pymongo.errors import AutoReconnect, NetworkTimeout, NotPrimaryError, ServerSelectionTimeoutError

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

_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: Dict[str, List[Any]] = 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 upsert(self, collection: str, filter: Dict[str, Any], data: Dict[str, Any]) -> bool:
    """
    Atomically creates or updates a single document matching the filter.

    Args:
        collection (str): The name of the collection.
        filter (Dict[str, Any]): Filter criteria to match a single document.
        data (Dict[str, Any]): Update operations to apply.

    Returns:
        bool: True if a new document was inserted, False if an existing one was updated.
    """
    self._init_index(collection)
    col = self.db.get_collection(collection)
    search_filter = self._convert_filter_id(filter)
    result = col.update_one(search_filter, data, upsert=True)
    return result.upserted_id is not None

  @_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

  @_mongo_retry
  @exception_logger
  def count(self, collection: str, filter: Optional[Dict[str, Any]] = None) -> int:
    """
    Counts documents in the specified collection matching the filter.

    Args:
        collection (str): The name of the collection to query.
        filter (Dict[str, Any], optional): Filter criteria. Defaults to no filter.

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

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