SeaCogs/aurora/models.py

import json
from datetime import datetime, timedelta
from typing import Dict, List, Optional

from pydantic import BaseModel


class Moderation(BaseModel):
    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

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

    @classmethod
    def from_sql(cls, moderation_id: int, guild_id: int):
        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:
                if result[7] is not None:
                    hours, minutes, seconds = map(int, result[7].split(':'))
                    duration = timedelta(hours=hours, minutes=minutes, seconds=seconds)
                else:
                    duration = None
                case = {
                    "moderation_id": int(result[0]),
                    "guild_id": int(guild_id),
                    "timestamp": datetime.fromtimestamp(result[1]),
                    "moderation_type": str(result[2]),
                    "target_type": str(result[3]),
                    "target_id": int(result[4]),
                    "moderator_id": int(result[5]),
                    "role_id": int(result[6]) if result[6] is not None else None,
                    "duration": duration,
                    "end_timestamp": datetime.fromtimestamp(result[8]) if result[8] is not None else None,
                    "reason": result[9],
                    "resolved": bool(result[10]),
                    "resolved_by": result[11],
                    "resolve_reason": result[12],
                    "expired": bool(result[13]),
                    "changes": result[14] if result[14] else [],
                    "metadata": result[15] if result[15] else {},
                }

                cursor.close()

                return cls(**case)

        return None

class JSONEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, datetime):
            return int(o.timestamp())
        if isinstance(o, timedelta):
            return str(o)
        if isinstance(o, Moderation):
            return o.model_dump()
        if isinstance(o, None):
            return "NULL"
        return super().default(o)