SeaCogs/aurora/models.py

from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Union

from async_property import async_cached_property
from discord import Forbidden, HTTPException, InvalidData, NotFound
from pydantic import BaseModel
from redbot.core.bot import Red

from aurora.utilities.utils import generate_dict


class AuroraBaseModel(BaseModel):
    """Base class for all models in Aurora."""
class Moderation(AuroraBaseModel):
    bot: Red
    moderation_id: int
    guild_id: int
    timestamp: datetime
    moderation_type: str
    target_type: str
    target_id: int
    moderator_id: int
    role_id: Optional[int]
    duration: Optional[timedelta]
    end_timestamp: Optional[datetime]
    reason: Optional[str]
    resolved: bool
    resolved_by: Optional[int]
    resolve_reason: Optional[str]
    expired: bool
    changes: List[Dict]
    metadata: Dict

    @property
    def id(self) -> int:
        return self.moderation_id

    @property
    def type(self) -> str:
        return self.moderation_type

    @async_cached_property
    async def moderator(self) -> "PartialUser":
        return await PartialUser.from_id(self.bot, self.moderator_id)

    @async_cached_property
    async def target(self) -> Union["PartialUser", "PartialChannel"]:
        if self.target_type == "user":
            return await PartialUser.from_id(self.bot, self.target_id)
        else:
            return await PartialChannel.from_id(self.bot, self.target_id)

    @async_cached_property
    async def resolved_by_user(self) -> Optional["PartialUser"]:
        if self.resolved_by:
            return await PartialUser.from_id(self.bot, self.resolved_by)
        return None

    def __str__(self):
        return f"{self.moderation_type} {self.target_type} {self.target_id} {self.reason}"

    @classmethod
    def from_sql(cls, bot: Red, moderation_id: int, guild_id: int) -> Optional["Moderation"]:
        from aurora.utilities.database import connect
        query = f"SELECT * FROM moderation_{guild_id} WHERE moderation_id = ?;"

        with connect() as database:
            cursor = database.cursor()
            cursor.execute(query, (moderation_id,))
            result = cursor.fetchone()

            if result:
                case = generate_dict(result)
                cursor.close()
                return cls.from_dict(bot, case)

        return None

    @classmethod
    def from_dict(cls, bot: Red, data: dict) -> "Moderation":
        return cls(bot=bot, **data)

    def to_json(self, indent: int = None, file: Any = None):
        from aurora.utilities.json import dump, dumps
        return dump(self.model_dump(exclude={"bot", "guild_id"}), file, indent=indent) if file else dumps(self.model_dump(exclude={"bot", "guild_id"}), indent=indent)


class PartialUser(AuroraBaseModel):
    id: int
    username: str
    discriminator: int

    @property
    def name(self):
        return f"{self.username}#{self.discriminator}" if self.discriminator == 0 else self.username

    def __str__(self):
        return self.name

    @classmethod
    async def from_id(cls, bot: Red, user_id: int) -> "PartialUser":
        user = bot.get_user(user_id)
        if not user:
            try:
                user = await bot.fetch_user(user_id)
                return cls(id=user.id, username=user.name, discriminator=user.discriminator)
            except NotFound:
                return cls(id=user_id, username="Deleted User", discriminator=0)

class PartialChannel(AuroraBaseModel):
    id: int
    name: str

    @property
    def mention(self):
        if self.name == "Deleted Channel" or self.name == "Forbidden Channel":
            return self.name
        return f"<#{self.id}>"

    def __str__(self):
        return self.mention

    @classmethod
    async def from_id(cls, bot: Red, channel_id: int) -> "PartialChannel":
        user = bot.get_channel(channel_id)
        if not user:
            try:
                user = await bot.fetch_channel(channel_id)
                return cls(id=user.id, username=user.name, discriminator=user.discriminator)
            except (NotFound, InvalidData, HTTPException, Forbidden) as e:
                if e == Forbidden:
                    return cls(id=channel_id, name="Forbidden Channel")
                return cls(id=channel_id, name="Deleted Channel")