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]) -> 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
        """
        if (
            await col_users.find_one(
                {"id": user.id if isinstance(user, User) else user}
            )
            is None
        ):
            await col_users.insert_one(
                {
                    "id": user.id if isinstance(user, User) else user,
                    "locale": user.language_code if isinstance(user, User) else None,
                }
            )

        db_record = await col_users.find_one(
            {"id": user.id if isinstance(user, User) else user}
        )

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

        return PyroUser(**db_record)