SeaCogs/aurora/models.py

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

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

from aurora.utilities.utils import generate_dict


class AuroraBaseModel(BaseModel):
    """Base class for all models in Aurora."""
    model_config = ConfigDict(ignored_types=(Red,), arbitrary_types_allowed=True)
    bot: Red

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

class AuroraGuildModel(AuroraBaseModel):
    """Subclass of AuroraBaseModel that includes a guild_id attribute, and a modified to_json() method to match."""
    guild_id: int

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

class Moderation(AuroraGuildModel):
    moderation_id: int
    timestamp: datetime
    moderation_type: str
    target_type: str
    target_id: int
    moderator_id: int
    role_id: Optional[int] = None
    duration: Optional[timedelta] = None
    end_timestamp: Optional[datetime] = None
    reason: Optional[str] = None
    resolved: bool
    resolved_by: Optional[int] = None
    resolve_reason: Optional[str] = None
    expired: bool
    changes: List["Change"]
    metadata: Dict

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

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

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

    async def get_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 def get_resolved_by(self) -> Optional["PartialUser"]:
        if self.resolved_by:
            return await PartialUser.from_id(self.bot, self.resolved_by)
        return None

    async def get_role(self) -> Optional["PartialRole"]:
        if self.role_id:
            return await PartialRole.from_id(self.bot, self.guild_id, self.role_id)
        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(bot, result, guild_id)
                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)

class Change(AuroraBaseModel):
    type: Literal["ORIGINAL", "RESOLVE", "EDIT"]
    timestamp: datetime
    reason: str
    user_id: int
    duration: Optional[timedelta] = None
    end_timestamp: Optional[datetime] = None

    def __str__(self):
        return f"{self.type} {self.user_id} {self.reason}"

    async def get_user(self) -> "PartialUser":
        return await PartialUser.from_id(self.bot, self.user_id)

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

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(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator)
            except NotFound:
                return cls(bot=bot, id=user_id, username="Deleted User", discriminator=0)


class PartialChannel(AuroraGuildModel):
    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":
        channel = bot.get_channel(channel_id)
        if not channel:
            try:
                channel = await bot.fetch_channel(channel_id)
                return cls(bot=bot, guild_id=channel.guild.id, id=channel.id, username=channel.name)
            except (NotFound, InvalidData, HTTPException, Forbidden) as e:
                if e == Forbidden:
                    return cls(bot=bot, guild_id=0, id=channel_id, name="Forbidden Channel")
                return cls(bot=bot, guild_id=0, id=channel_id, name="Deleted Channel")

class PartialRole(AuroraGuildModel):
    id: int
    name: str

    @property
    def mention(self):
        return f"<@&{self.id}>"

    def __str__(self):
        return self.mention

    @classmethod
    async def from_id(cls, bot: Red, guild_id: int, role_id: int) -> "PartialRole":
        try:
            guild = await bot.fetch_guild(guild_id, with_counts=False)
        except (Forbidden, HTTPException):
            return cls(bot=bot, guild_id=guild_id, id=role_id, name="Forbidden Role")
        role = guild.get_role(role_id)
        if not role:
            return cls(bot=bot, guild_id=guild_id, id=role_id, name="Deleted Role")
        return cls(bot=bot, guild_id=guild_id, id=role.id, name=role.name)