Vocard/voicelink/mongodb.py

"""MIT License

Copyright (c) 2023 - present Vocard Development

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import copy
import time
import asyncio
import logging

from typing import Any, Dict, Optional, Literal, TypedDict, List
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection

from .config import Config

logger: logging.Logger = logging.getLogger("vocard.db")

# Type definitions for better code clarity
class PlaylistPerms(TypedDict):
    read: List[int]
    write: List[int]
    remove: List[int]

class Playlist(TypedDict):
    tracks: List[Dict[str, Any]]
    perms: PlaylistPerms
    name: str
    type: Literal["playlist"]

class UserData(TypedDict):
    _id: int
    playlist: Dict[str, Playlist]
    history: List[Dict[str, Any]]
    inbox: List[Dict[str, Any]]

UpdateOperationType = Literal["$set", "$unset", "$inc", "$push", "$pull"]


class MongoDBHandler:
    """
    Handles MongoDB operations with connection pooling and caching.
    Implements a thread-safe singleton pattern for database connections.
    """
    
    # Static instance variables
    _client: Optional[AsyncIOMotorClient] = None
    _db: Optional[Any] = None
    _settings_db: Optional[AsyncIOMotorCollection] = None
    _users_db: Optional[AsyncIOMotorCollection] = None
    _lock: asyncio.Lock = asyncio.Lock()

    # Cache with TTL (Time To Live in seconds)
    _CACHE_TTL: int = 300  # 5 minutes
    _settings_buffer: Dict[int, Dict[str, Any]] = {}
    _users_buffer: Dict[int, Dict[str, Any]] = {}
    _last_access: Dict[int, float] = {}  # Tracks last access time for cache entries

    # Maximum cache size to prevent memory issues
    _MAX_CACHE_SIZE: int = 10000

    # Default user template
    _user_base: UserData = {
        "_id": 0,  # Will be replaced with actual user ID
        "playlist": {
            "200": {
                "tracks": [],
                "perms": {"read": [], "write": [], "remove": []},
                "name": "Favourite",
                "type": "playlist",
            }
        },
        "history": [],
        "inbox": [],
    }

    @classmethod
    async def init(cls, uri: str, db_name: str) -> None:
        """
        Initialize the MongoDB connection with connection pooling and error handling.
        
        Args:
            uri (str): MongoDB connection URI
            db_name (str): Name of the database to use
            
        Raises:
            ConnectionError: If unable to connect to MongoDB
            Exception: For other initialization errors
        """
        if not uri or not db_name:
            logger.error("MongoDB initialization failed: URI or database name is missing.")
            raise ValueError("Both URI and database name must be provided.")

        async with cls._lock:
            if cls._client is not None:
                logger.warning("MongoDB client is already initialized. Skipping reinitialization.")
                return

            logger.debug("Initializing MongoDB client with URI: %s and DB name: %s", uri, db_name)

            try:
                cls._client = AsyncIOMotorClient(
                    uri,
                    maxPoolSize=50,
                    minPoolSize=5,
                    maxIdleTimeMS=60000,
                    retryWrites=True
                )
                logger.debug("MongoDB client created successfully. Testing connection...")

                await cls._client.server_info()
                logger.debug("MongoDB connection test passed.")

                cls._db = cls._client[db_name]
                cls._settings_db = cls._db["Settings"]
                cls._users_db = cls._db["Users"]

                logger.info("MongoDB databases initialized: %s", db_name)

            except Exception as e:
                logger.error("MongoDB initialization failed: %s", str(e), exc_info=True)

                cls._client = None
                cls._db = None
                cls._settings_db = None
                cls._users_db = None

                raise ConnectionError(f"Failed to initialize MongoDB: {str(e)}")
                
    @classmethod
    async def cleanup_cache(cls) -> None:
        """
        Cleanup expired cache entries to prevent memory leaks.
        Should be called periodically or when cache size exceeds _MAX_CACHE_SIZE.
        """
        current_time = time.time()
        logger.info("Starting cache cleanup at timestamp: %.2f", current_time)

        async with cls._lock:
            try:
                # Remove expired entries from settings cache
                expired_settings = [
                    guild_id for guild_id, last_access in cls._last_access.items()
                    if current_time - last_access > cls._CACHE_TTL and guild_id in cls._settings_buffer
                ]
                logger.debug("Found %d expired cache entries.", len(expired_settings))

                for guild_id in expired_settings:
                    del cls._settings_buffer[guild_id]
                    del cls._last_access[guild_id]
                    logger.debug("Removed expired cache for guild_id: %s", guild_id)

                # If still too large, remove oldest entries
                while len(cls._settings_buffer) > cls._MAX_CACHE_SIZE:
                    oldest_id = min(cls._last_access.items(), key=lambda x: x[1])[0]
                    del cls._settings_buffer[oldest_id]
                    del cls._last_access[oldest_id]
                    logger.warning("Cache size exceeded. Removed oldest entry: %s", oldest_id)

                logger.info("Cache cleanup completed. Current cache size: %d", len(cls._settings_buffer))

            except Exception as e:
                logger.error("Cache cleanup failed: %s", str(e), exc_info=True)

    @classmethod
    async def _update_db(
        cls,
        db: AsyncIOMotorCollection,
        cache: Dict[str, Any],
        filter_: Dict[str, Any],
        data: Dict[UpdateOperationType, Dict[str, Any]],
    ) -> bool:
        """
        Update database and cache atomically with error handling and validation.
        
        Args:
            db: MongoDB collection to update
            cache: Cache dictionary to update
            filter_: MongoDB filter for the update
            data: Update operations to perform
            
        Returns:
            bool: True if update was successful, False otherwise
            
        Raises:
            ValueError: If invalid update operation is provided
        """
        async with cls._lock:
            try:
                # Validate update operations
                valid_operations = {"$set", "$unset", "$inc", "$push", "$pull"}
                if not all(op in valid_operations for op in data.keys()):
                    raise ValueError(f"Invalid update operation. Must be one of {valid_operations}")

                # Update cache first
                for mode, action in data.items():
                    for key, value in action.items():
                        cursors = key.split(".")
                        nested = cache
                        
                        # Ensure path exists
                        for c in cursors[:-1]:
                            if not isinstance(nested, dict):
                                raise ValueError(f"Invalid path: {key}")
                            nested = nested.setdefault(c, {})

                        field = cursors[-1]

                        try:
                            if mode == "$set":
                                nested[field] = value
                            elif mode == "$unset":
                                nested.pop(field, None)
                            elif mode == "$inc":
                                if not isinstance(nested.get(field, 0), (int, float)):
                                    raise ValueError(f"Cannot increment non-numeric field: {field}")
                                nested[field] = nested.get(field, 0) + value
                            elif mode == "$push":
                                arr = nested.setdefault(field, [])
                                if not isinstance(arr, list):
                                    raise ValueError(f"Cannot push to non-array field: {field}")
                                if isinstance(value, dict) and "$each" in value:
                                    arr.extend(value["$each"])
                                    if "$slice" in value:
                                        arr[:] = arr[value["$slice"]:]
                                else:
                                    arr.append(value)
                            elif mode == "$pull":
                                if field in nested:
                                    if not isinstance(nested[field], list):
                                        raise ValueError(f"Cannot pull from non-array field: {field}")
                                    values = value.get("$in", []) if isinstance(value, dict) else [value]
                                    nested[field] = [item for item in nested[field] if item not in values]
                        except Exception as e:
                            raise ValueError(f"Error updating {key}: {str(e)}")

                # Then update database
                result = await db.update_one(filter_, data)
                
                # Update last access time
                if '_id' in filter_:
                    cls._last_access[filter_['_id']] = time.time()
                
                return result.modified_count > 0

            except Exception as e:
                # Rollback cache if database update fails
                if '_id' in filter_:
                    cls._settings_buffer.pop(filter_['_id'], None)
                    cls._users_buffer.pop(filter_['_id'], None)
                raise Exception(f"Update failed: {str(e)}")

    @classmethod
    def get_cached_settings(
        cls, 
        guild_id: int
    ) -> Dict[str, Any]:
        """
        Retrieve settings for a guild with caching.
        
        Args:
            guild_id: The Discord guild ID
            
        Returns:
            Dict containing guild settings or empty dict if not found
        """
        try:
            if guild_id not in cls._settings_buffer:
                return {}

            return copy.deepcopy(cls._settings_buffer[guild_id])

        except Exception as e:
            raise ConnectionError(f"Failed to retrieve settings: {str(e)}")
        
    @classmethod
    async def get_settings(
        cls, 
        guild_id: int,
        *,
        deep_copy: bool = True,
        force_refresh: bool = False,
    ) -> Dict[str, Any]:
        """
        Retrieve settings for a guild with caching.
        
        Args:
            guild_id: The Discord guild ID
            force_refresh: If True, bypass cache and fetch fresh data
            
        Returns:
            Dict containing guild settings
            
        Raises:
            ConnectionError: If database operation fails
        """
        try:
            async with cls._lock:
                # Check if we need fresh data
                if force_refresh or guild_id not in cls._settings_buffer:
                    settings = await cls._settings_db.find_one({"_id": guild_id})
                    if not settings:
                        settings = {"_id": guild_id}
                        try:
                            await cls._settings_db.insert_one(settings)
                        except Exception as e:
                            raise ConnectionError(f"Failed to create settings: {str(e)}")
                    
                    cls._settings_buffer[guild_id] = settings
                    cls._last_access[guild_id] = time.time()
                
                buffer = cls._settings_buffer[guild_id]
                return copy.deepcopy(buffer) if deep_copy else buffer

        except Exception as e:
            raise ConnectionError(f"Failed to retrieve settings: {str(e)}")

    @classmethod
    async def update_settings(
        cls,
        guild_id: int,
        data: Dict[UpdateOperationType, Dict[str, Any]],
        *,
        upsert: bool = False
    ) -> bool:
        """
        Update settings for a guild.
        
        Args:
            guild_id: The Discord guild ID
            data: Update operations to perform
            upsert: If True, create document if it doesn't exist
            
        Returns:
            bool: True if update was successful
            
        Raises:
            ValueError: If invalid update data is provided
            ConnectionError: If database operation fails
        """
        try:
            settings = await cls.get_settings(guild_id, deep_copy=False)
            result = await cls._update_db(
                cls._settings_db,
                settings,
                {"_id": guild_id},
                data
            )
            
            if not result and upsert:
                # Try to insert if update failed and upsert is True
                settings = {"_id": guild_id, **data.get("$set", {})}
                await cls._settings_db.insert_one(settings)
                cls._settings_buffer[guild_id] = settings
                return True
                
            return result
            
        except Exception as e:
            raise ConnectionError(f"Failed to update settings: {str(e)}")

    @classmethod
    async def get_user(
        cls,
        user_id: int,
        *,
        d_type: Optional[str] = None,
        need_copy: bool = True,
        force_refresh: bool = False
    ) -> Dict[str, Any]:
        """
        Retrieve user data with caching and type-specific data.
        
        Args:
            user_id: The Discord user ID
            d_type: Specific data type to retrieve
            need_copy: If True, return a deep copy of the data
            force_refresh: If True, bypass cache and fetch fresh data
            
        Returns:
            Dict containing user data
            
        Raises:
            ConnectionError: If database operation fails
            ValueError: If invalid d_type is provided
        """
        try:
            async with cls._lock:
                # Check if we need fresh data
                if force_refresh or user_id not in cls._users_buffer:
                    user = await cls._users_db.find_one({"_id": user_id})
                    if not user:
                        user = {**copy.deepcopy(cls._user_base), "_id": user_id}
                        user["playlist"]["200"]["name"] = Config().get_playlist_config()[2]
                        try:
                            await cls._users_db.insert_one(user)
                        except Exception as e:
                            raise ConnectionError(f"Failed to create user: {str(e)}")
                    
                    cls._users_buffer[user_id] = user
                    cls._last_access[user_id] = time.time()

                user = cls._users_buffer[user_id]
                
                if d_type:
                    if d_type not in cls._user_base:
                        raise ValueError(f"Invalid data type: {d_type}")
                    user = user.setdefault(d_type, copy.deepcopy(cls._user_base.get(d_type)))

                return copy.deepcopy(user) if need_copy else user

        except Exception as e:
            raise ConnectionError(f"Failed to retrieve user data: {str(e)}")

    @classmethod
    async def update_user(
        cls,
        user_id: int,
        data: Dict[UpdateOperationType, Dict[str, Any]],
        *,
        upsert: bool = False
    ) -> bool:
        """
        Update user data.
        
        Args:
            user_id: The Discord user ID
            data: Update operations to perform
            upsert: If True, create user if doesn't exist
            
        Returns:
            bool: True if update was successful
            
        Raises:
            ValueError: If invalid update data is provided
            ConnectionError: If database operation fails
        """
        try:
            user = await cls.get_user(user_id, need_copy=False)
            result = await cls._update_db(
                cls._users_db,
                user,
                {"_id": user_id},
                data
            )
            
            if not result and upsert:
                # Try to insert if update failed and upsert is True
                user_data = {"_id": user_id, **data.get("$set", {})}
                await cls._users_db.insert_one(user_data)
                cls._users_buffer[user_id] = user_data
                return True
                
            return result
            
        except Exception as e:
            raise ConnectionError(f"Failed to update user: {str(e)}")

    @classmethod
    async def delete_user(cls, user_id: int) -> bool:
        """
        Delete a user's data completely.
        
        Args:
            user_id: The Discord user ID
            
        Returns:
            bool: True if deletion was successful
            
        Raises:
            ConnectionError: If database operation fails
        """
        try:
            async with cls._lock:
                result = await cls._users_db.delete_one({"_id": user_id})
                if result.deleted_count > 0:
                    cls._users_buffer.pop(user_id, None)
                    cls._last_access.pop(user_id, None)
                    return True
                return False
                
        except Exception as e:
            raise ConnectionError(f"Failed to delete user: {str(e)}")

    @classmethod
    async def get_users_by_criteria(
        cls,
        criteria: Dict[str, Any],
        *,
        limit: Optional[int] = None,
        skip: int = 0
    ) -> List[Dict[str, Any]]:
        """
        Retrieve multiple users matching specific criteria.
        
        Args:
            criteria: MongoDB query criteria
            limit: Maximum number of users to return
            skip: Number of matching users to skip
            
        Returns:
            List of matching user data
            
        Raises:
            ConnectionError: If database operation fails
        """
        try:
            cursor = cls._users_db.find(criteria).skip(skip)
            if limit:
                cursor = cursor.limit(limit)
                
            users = await cursor.to_list(length=None)
            
            # Update cache with fetched users
            async with cls._lock:
                current_time = time.time()
                for user in users:
                    user_id = user["_id"]
                    cls._users_buffer[user_id] = user
                    cls._last_access[user_id] = current_time
                    
            return users
            
        except Exception as e:
            raise ConnectionError(f"Failed to retrieve users: {str(e)}")