YandexMusicDiscordBot/MusicBot/database/base.py

"""This documents initialises databse and contains methods to access it."""

from typing import Any, cast

from pymongo import MongoClient
from pymongo.collection import Collection

from .user import User, ExplicitUser
from .guild import Guild, ExplicitGuild, MessageVotes

client: MongoClient = MongoClient("mongodb://localhost:27017/")
users: Collection[ExplicitUser] = client.YandexMusicBot.users
guilds: Collection[ExplicitGuild] = client.YandexMusicBot.guilds

class BaseUsersDatabase:

    def create_record(self, uid: int) -> None:
        """Create user database record.

        Args:
            uid (int): User id.
        """
        uid = uid
        users.insert_one(ExplicitUser(
            _id=uid,
            ym_token=None,
            playlists=[],
            playlists_page=0,
            queue_page=0,
            vibe_batch_id=None,
            vibe_type=None,
            vibe_id=None,
            vibe_settings={
                'mood': 'all',
                'diversity': 'default',
                'lang': 'any'
            }
        ))

    def update(self, uid: int, data: User | dict[Any, Any]) -> None:
        """Update user record.

        Args:
            uid (int): User id.
            data (User | dict[Any, Any]): Updated data.
        """
        self.get_user(uid)
        users.update_one({'_id': uid}, {"$set": data})

    def get_user(self, uid: int) -> ExplicitUser:
        """Get user record from database. Create new entry if not present.

        Args:
            uid (int): User id.

        Returns:
            User: User record.
        """
        user = users.find_one({'_id': uid})
        if not user:
            self.create_record(uid)
            user = users.find_one({'_id': uid})
        user =  cast(ExplicitUser, user)
        existing_fields = user.keys()
        fields: ExplicitUser = ExplicitUser(
            _id=0,
            ym_token=None,
            playlists=[],
            playlists_page=0,
            queue_page=0,
            vibe_batch_id=None,
            vibe_type=None,
            vibe_id=None,
            vibe_settings={
                'mood': 'all',
                'diversity': 'default',
                'lang': 'any'
            }
        )
        for field, default_value in fields.items():
            if field not in existing_fields:
                user[field] = default_value
                users.update_one({'_id': uid}, {"$set": {field: default_value}})
        
        return user

    def get_ym_token(self, uid: int) -> str | None:
        user = users.find_one({'_id': uid})
        if not user:
            self.create_record(uid)
            user = users.find_one({'_id': uid})
        return cast(ExplicitUser, user)['ym_token']

class BaseGuildsDatabase:
    
    def create_record(self, gid: int) -> None:
        """Create guild database record.

        Args:
            gid (int): Guild id.
        """
        guilds.insert_one(ExplicitGuild(
            _id=gid,
            next_tracks=[],
            previous_tracks=[],
            current_track=None,
            current_menu=None,
            is_stopped=True,
            allow_explicit=True,
            always_allow_menu=False,
            vote_next_track=True,
            vote_add_track=True,
            vote_add_album=True,
            vote_add_artist=True,
            vote_add_playlist=True,
            shuffle=False,
            repeat=False,
            votes={},
            vibing=False,
            current_viber_id=None
        ))

    def update(self, gid: int, data: Guild) -> None:
        """Update guild record.

        Args:
            gid (int): Guild id.
            data (dict[Any, Any]): Updated data.
        """
        self.get_guild(gid)
        guilds.update_one({'_id': gid}, {"$set": data})

    def get_guild(self, gid: int) -> ExplicitGuild:
        """Get guild record from database. Create new entry if not present.

        Args:
            uid (int): User id.

        Returns:
            Guild: Guild record.
        """
        guild = guilds.find_one({'_id': gid})
        if not guild:
            self.create_record(gid)
            guild = guilds.find_one({'_id': gid})
        
        guild = cast(ExplicitGuild, guild)
        existing_fields = guild.keys()
        fields = ExplicitGuild(
            _id=0,
            next_tracks=[],
            previous_tracks=[],
            current_track=None,
            current_menu=None,
            is_stopped=True,
            allow_explicit=True,
            always_allow_menu=False,
            vote_next_track=True,
            vote_add_track=True,
            vote_add_album=True,
            vote_add_artist=True,
            vote_add_playlist=True,
            shuffle=False,
            repeat=False,
            votes={},
            vibing=False,
            current_viber_id=None
        )
        for field, default_value in fields.items():
            if field not in existing_fields:
                guild[field] = default_value
                guilds.update_one({'_id': gid}, {"$set": {field: default_value}})
        
        return guild
        
    def update_vote(self, gid: int, mid: int, data: MessageVotes) -> None:
        """Update vote for a message in a guild.
        
        Args:
            gid (int): Guild id.
            mid (int): Message id.
            vote (bool): Vote value.
        """
        guild = self.get_guild(gid)
        guild['votes'][str(mid)] = data
        guilds.update_one({'_id': gid}, {"$set": {'votes': guild['votes']}})