from typing import Union

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

# PyroClient uses MongoDB implementation of PyroUser
# but you can also select SQLite one below
# from classes.pyrouser_sqlite import PyroUser
from classes.pyrouser_mongo import PyroUser


class PyroClient(LibPyroClient):
    async def find_user(self, user: Union[int, User]) -> PyroUser:
        """Find User by its ID or User object.

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

        ### Returns:
            * `PyroUser`: User in database representation.
        """

        return (
            await PyroUser.find(user)
            if isinstance(user, int)
            else await PyroUser.find(user.id, locale=user.language_code)
        )