TelegramBot/classes/pyrouser.py

import logging
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any, Mapping, Tuple, Union

import pytz
from bson import ObjectId

from classes.location import Location
from modules.database import col_users

logger = logging.getLogger(__name__)


@dataclass
class PyroUser:
    """Dataclass of DB entry of a user"""

    __slots__ = (
        "_id",
        "id",
        "locale",
        "enabled",
        "location",
        "offset",
        "time_hour",
        "time_minute",
    )

    _id: ObjectId
    id: int
    locale: Union[str, None]
    enabled: bool
    location: Union[Location, None]
    offset: int
    time_hour: int
    time_minute: int

    @classmethod
    async def find(
        cls,
        id: int,
        locale: Union[str, None] = None,
        enabled: bool = True,
        location_id: int = 0,
        offset: int = 1,
        time_hour: int = 16,
        time_minute: int = 0,
    ) -> "PyroUser":
        db_entry = await col_users.find_one({"id": id})

        if db_entry is None:
            inserted = await col_users.insert_one(
                {
                    "id": id,
                    "locale": locale,
                    "enabled": enabled,
                    "location": location_id,
                    "offset": offset,
                    "time_hour": time_hour,
                    "time_minute": time_minute,
                }
            )
            db_entry = await col_users.find_one({"_id": inserted.inserted_id})

        if db_entry is None:
            raise RuntimeError("Could not find inserted user entry.")

        try:
            db_entry["location"] = await Location.get(db_entry["location"])  # type: ignore
        except ValueError:
            db_entry["location"] = None  # type: ignore

        return cls(**db_entry)

    @classmethod
    async def from_dict(cls, **kwargs) -> "PyroUser":
        if "location" in kwargs:
            try:
                kwargs["location"] = await Location.get(kwargs["location"])  # type: ignore
            except ValueError:
                kwargs["location"] = None  # type: ignore
        return cls(**kwargs)

    async def update_locale(self, locale: Union[str, None]) -> Union[str, None]:
        """Change user's locale stored in the database.

        ### Args:
            * locale (`Union[str, None]`): New locale to be set.
        """
        logger.info("%s's locale has been set to %s", self.id, locale)

        await col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})

        self.locale = locale

        return self.locale

    async def update_state(self, enabled: bool = False) -> bool:
        """Update user's state (enabled/disabled)

        ### Args:
            * enabled (`bool`, *optional*): Whether the user is enabled. Defaults to `False`.

        ### Returns:
            * `bool`: User's current state
        """
        logger.info("%s's state has been set to %s", self.id, enabled)

        await col_users.update_one({"_id": self._id}, {"$set": {"enabled": enabled}})

        self.enabled = enabled

        return self.enabled

    async def update_location(self, location_id: int) -> Location:
        """Update user's location and move their time to the new timezone (if the user had a location set previously)

        ### Args:
            * location_id (`int`): ID of the location

        ### Returns:
            `Location`: New location
        """
        logger.info("%s's location has been set to %s", self.id, location_id)

        await col_users.update_one(
            {"_id": self._id}, {"$set": {"location": location_id}}
        )

        location = await Location.get(location_id)

        # Execute if timezones of old and new locations are different
        if self.location and (self.location.timezone.zone != location.timezone.zone):
            # Get UTC time for selected reminder time
            now_utc = datetime.now(pytz.utc).replace(
                hour=self.time_hour, minute=self.time_minute, second=0, microsecond=0
            )

            # Get the time for the reminder time of old and new location
            local_old = now_utc.astimezone(self.location.timezone)
            local_new = (
                location.timezone.localize(local_old.replace(tzinfo=None))
            ).astimezone(pytz.utc)

            # Update the time to match the new timezone
            await self.update_time(hour=local_new.hour, minute=local_new.minute)

        self.location = location

        return self.location

    async def update_offset(self, offset: int = 1) -> int:
        """Update the offset of the reminder (in days)

        ### Args:
            * offset (`int`, *optional*): Offset in days. Defaults to `1`.

        ### Returns:
            * `int`: Offset in days
        """
        logger.info("%s's offset has been set to %s", self.id, offset)

        await col_users.update_one({"_id": self._id}, {"$set": {"offset": offset}})

        self.offset = offset

        return offset

    async def update_time(self, hour: int = 16, minute: int = 0) -> Tuple[int, int]:
        """Update the time of the reminder (hour and minute, for UTC timezone)

        ### Args:
            * hour (`int`, *optional*): Hour of the reminder. Defaults to `16`.
            * minute (`int`, *optional*): Minute of the reminder. Defaults to `0`.

        ### Returns:
            * `Tuple[int, int]`: Hour and minute of the reminder
        """
        logger.info("%s's time has been set to %s h. %s m.", self.id, hour, minute)

        await col_users.update_one(
            {"_id": self._id}, {"$set": {"time_hour": hour, "time_minute": minute}}
        )

        self.time_hour = hour
        self.time_minute = minute

        return self.time_hour, self.time_minute

    async def delete(self) -> None:
        """Delete the database record of the user"""
        logger.info("%s's data has been deleted", self.id)

        await col_users.delete_one({"_id": self._id})

    async def checkout(self) -> Mapping[str, Any]:
        """Checkout the user's database record

        ### Raises:
            * `KeyError`: Database record of the user was not found

        ### Returns:
            * `Mapping[str, Any]`: Database record
        """
        logger.info("%s's data has been checked out", self.id)

        db_entry = await col_users.find_one({"_id": self._id})

        if db_entry is None:
            raise KeyError(
                f"DB record with id {self._id} of user {self.id} was not found"
            )

        del db_entry["_id"]  # type: ignore

        return db_entry

    def get_reminder_date(self) -> datetime:
        """Get next reminder date (year, month and day)

        ### Raises:
            * `AttributeError`: Some attribute(s) are missing

        ### Returns:
            * `datetime`: Datetime object of the next reminder date
        """
        if self.location is None:
            logger.warning(
                "Could not get the reminder date for %s: User does not have some attribute(s) set",
                self.id,
            )
            raise AttributeError(
                f"Could not get the reminder date for {self.id}: User does not have some attribute(s) set"
            )

        if not self.location.timezone:
            logger.warning("Location %s does not have a timezone set", self.location.id)

        return (
            datetime.now(self.location.timezone or pytz.utc) + timedelta(days=1)
        ).replace(hour=0, minute=0, second=0, microsecond=0)

    def get_reminder_time(self) -> datetime:
        """Get reminder time (hour and minute)

        ### Raises:
            * `AttributeError`: Some attribute(s) are missing

        ### Returns:
            * `datetime`: Datetime object of the next reminder date
        """
        if self.time_hour is None or self.time_minute is None or self.location is None:
            logger.warning(
                "Could not get the reminder time for %s: User does not have some attribute(s) set",
                self.id,
            )
            raise AttributeError(
                f"Could not get the reminder time for {self.id}: User does not have some attribute(s) set"
            )

        if not self.location.timezone:
            logger.warning("Location %s does not have a timezone set", self.location.id)

        return (
            datetime.now(pytz.utc)
            .replace(
                hour=self.time_hour,
                minute=self.time_minute,
                second=0,
                microsecond=0,
            )
            .astimezone(self.location.timezone or pytz.utc)
        )