from typing import Union

from libbot.pyrogram.classes import PyroClient
from pyrogram.types import User

from classes.pyrouser import PyroUser
from modules.database import col_users


class PyroClient(PyroClient):
    async def find_user(self, user: Union[int, User], group: int) -> PyroUser:
        """Find User by it's ID or User object

        ### Args:
            * user (`Union[int, User]`): ID or User object to extract ID from
            * group (`int`): ID of the group

        ### Returns:
            * `PyroUser`: PyroUser object
        """
        db_record = col_users.find_one(
            {"id": user.id if isinstance(user, User) else user, "group": group}
        )

        if db_record is None:
            raise KeyError(
                f"User with ID {user.id if isinstance(user, User) else user} was not found in the database"
            )

        return PyroUser(**db_record)