payphone/payphone/db/models.py

from datetime import datetime, timezone

import enum

from tortoise import models, fields

import redis.asyncio as redis

import orjson


class Report(models.Model):
    """
    A report of a message.
    """

    id = fields.BigIntField(pk=True)

    message_id = fields.BigIntField(unique=True)
    message_content = fields.TextField()
    channel_id = fields.BigIntField()
    user_id = fields.BigIntField()
    user_name = fields.TextField()
    user_avatar_url = fields.TextField()
    reply_message_content = fields.TextField(null=True, default=None)
    reply_user_id = fields.BigIntField(null=True)
    reply_user_name = fields.TextField(null=True, default=None)
    reply_user_avatar_url = fields.TextField(null=True, default=None)
    created_at = fields.DatetimeField(auto_now_add=True)
    resolved = fields.BooleanField(default=False, index=True)
    reported_by = fields.ForeignKeyField(
        "models.User", related_name="reports", null=True
    )
    resolved_at = fields.DatetimeField(null=True)
    resolved_by = fields.ForeignKeyField(
        "models.User", related_name="resolved_reports", null=True
    )


class UserFlag(enum.IntFlag):
    """
    Flags for a user.
    """

    NONE = 0
    MODERATOR = 1
    SUPPORTER = 2


class User(models.Model):
    """
    A user.
    """

    id = fields.BigIntField(pk=True)

    warns_count: int = fields.IntField(default=0)
    bans_count: int = fields.IntField(default=0)

    is_masked = fields.BooleanField(default=False)
    is_moderator = fields.BooleanField(default=False)
    is_supporter_until = fields.DatetimeField(null=True)

    @property
    def flag(self) -> UserFlag:
        """
        The flags of the user.
        """

        flags = UserFlag.NONE

        if self.is_moderator:
            flags |= UserFlag.MODERATOR

        if (
            self.is_supporter_until is not None
            and self.is_supporter_until > datetime.now(timezone.utc)
        ):
            flags |= UserFlag.SUPPORTER

        return flags

    @classmethod
    async def get_cache_or_db(cls, redis_client: redis.Redis, user_id: int) -> "User":
        """
        Returns the user from the cache or the database.

        Args:
            redis_client (redis.Redis): The Redis client.
            user_id (int): The user ID.

        Returns:
            (User): The user.
        """

        cache = await redis_client.get(f"user:{user_id}")

        instance = User(id=user_id)

        if cache is not None:
            cache = orjson.loads(cache)

            instance.warns_count = cache["warns_count"]
            instance.bans_count = cache["bans_count"]
            instance.is_masked = cache["is_masked"]
            instance.is_moderator = cache["is_moderator"]
            instance.is_supporter_until = (
                datetime.fromtimestamp(cache["is_supporter_until"], timezone.utc)
                if cache["is_supporter_until"] is not None
                else None
            )

        else:
            instance, _ = await User.update_or_create(id=user_id)

            await instance.set_cache(redis_client)

        return instance

    async def set_cache(self, redis_client: redis.Redis) -> None:
        """
        Sets the user in the cache.

        Args:
            redis_client (redis.Redis): The Redis client.
        """

        await redis_client.set(
            f"user:{self.id}",
            orjson.dumps(
                {
                    "warns_count": self.warns_count,
                    "bans_count": self.bans_count,
                    "is_masked": self.is_masked,
                    "is_moderator": self.is_moderator,
                    "is_supporter_until": (
                        self.is_supporter_until.timestamp()
                        if self.is_supporter_until
                        else None
                    ),
                }
            ),
            ex=60 * 15,
        )


class Warn(models.Model):
    """A warn."""

    user_id = fields.BigIntField()
    report = fields.ForeignKeyField("models.Report", null=True)
    reason = fields.TextField(null=True, default=None)
    created_at = fields.DatetimeField(auto_now_add=True)
    created_by = fields.ForeignKeyField("models.User", related_name="created_warns")


class Ban(models.Model):
    """
    A ban.
    """

    user_id = fields.BigIntField()
    report = fields.ForeignKeyField("models.Report", null=True)
    reason = fields.TextField(null=True, default=None)
    created_at = fields.DatetimeField(auto_now_add=True)
    expires_at = fields.DatetimeField(null=True)
    created_by = fields.ForeignKeyField("models.User", related_name="created_bans")
    expired_by = fields.ForeignKeyField(
        "models.User", related_name="expired_bans", null=True
    )

    @property
    def is_expired(self) -> bool:
        """
        Whether the ban is expired.
        """

        return self.expires_at is not None and self.expires_at < datetime.now(
            timezone.utc
        )

    class Meta:
        """
        Meta information.
        """

        indexes = [
            models.Index(fields=["user_id", "expires_at"]),
        ]


class Channel(models.Model):
    """
    A channel.
    """

    id = fields.BigIntField(pk=True)
    server_id = fields.BigIntField(index=True)

    is_enabled = fields.BooleanField(null=True)

    class Meta:
        """
        Meta information.
        """

        indexes = [
            models.Index(fields=["server_id", "is_enabled"]),
        ]


class Server(models.Model):
    """
    A server.
    """

    id = fields.BigIntField(pk=True)
    name = fields.TextField()
    prefix = fields.TextField(default="p.")
    language = fields.TextField(default="en")

    async def set_prefix(self, redis_client: redis.Redis, prefix: str) -> None:
        """
        Sets the prefix of the server.

        Args:
            redis_client (redis.Redis): The Redis client.
            prefix (str): The prefix.
        """

        self.prefix = prefix

        await self.save()

        await redis_client.set(f"prefix:{self.id}", prefix, ex=60 * 60 * 24)