# filename: database/connection.py

import logging
from motor.motor_asyncio import AsyncIOMotorClient
import config

logger = logging.getLogger(__name__)

class Database:
    """
    Manages connections to one or more MongoDB databases.
    """
    def __init__(self, uris: list[str]):
        if not uris:
            raise ValueError("At least one MongoDB URI is required.")
            
        self.clients = [AsyncIOMotorClient(uri) for uri in uris]
        self.databases = [client.get_default_database() for client in self.clients]
        
        # For data that should not be split (like cache), we use the first DB as the primary.
        self.primary_db = self.databases[0]
        
        # For distributing new users, we'll cycle through the available databases.
        self._user_db_round_robin_counter = 0
        
        logger.info(f"Successfully connected to {len(self.clients)} MongoDB database(s).")

    async def get_user_db(self, user_id: int):
        """
        Gets the database assigned to a specific user.
        This ensures a user's data always stays in the same database.
        We use a simple hashing method to distribute users.
        """
        db_index = user_id % len(self.databases)
        return self.databases[db_index]
        
    async def find_user_db(self, user_id: int):
        """
        Searches across all databases to find which one contains the user.
        Returns the database object if found, otherwise None.
        """
        for db in self.databases:
            if await db.users.find_one({"user_id": user_id}):
                return db
        return None

    def get_next_db_for_new_user(self):
        """
        Picks a database for a new user using a round-robin strategy.
        This helps distribute new users evenly.
        """
        db = self.databases[self._user_db_round_robin_counter]
        self._user_db_round_robin_counter = (self._user_db_round_robin_counter + 1) % len(self.databases)
        return db

# --- Initialize the Database Connection ---
# This single `db` object will be imported and used by the rest of the application.
db = Database(config.MONGODB_URIS)