from typing import Union

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

from classes.pyrogroup import PyroGroup
from classes.pyrouser import PyroUser


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

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

        ### Returns:
            * `PyroUser`: PyroUser object
        """

        return (
            await PyroUser.find(user, *args, **kwargs)
            if isinstance(user, int)
            else await PyroUser.find(user.id, *args, **kwargs)
        )

    async def find_group(self, id: int, *args, **kwargs) -> PyroGroup:
        """Find Group by it's ID.

        ### Args:
            * id (`int`): Group ID

        ### Returns:
            * `PyroGroup`: PyroGroup object
        """

        return await PyroGroup.find(id, *args, **kwargs)