payphone/payphone/cogs/admin/cog.py

from datetime import datetime, timezone

import discord
from discord import errors
from discord.ext import commands, tasks

from tortoise import transactions, functions, exceptions

from payphone import bot
from payphone.db import utils as dbutils, models
from payphone.cogs import base
from payphone.core import utils


class Admin(base.Cog):
    def __init__(self, bot_instance: bot.Payphone):
        self.bot = bot_instance

        self.update_leaderboard.add_exception_type(exceptions.ConfigurationError)
        self.update_leaderboard.start()

    def cog_unload(self) -> None:
        self.update_leaderboard.cancel()

    @tasks.loop(seconds=30)
    async def update_leaderboard(self) -> None:
        """Updates the leaderboard."""

        calls = (
            await models.Call.annotate(messages_count=functions.Count("messages"))
            .filter(
                started_at__month=datetime.now(timezone.utc).month,
                started_at__year=datetime.now(timezone.utc).year,
                messages_count__gte=25,
            )
            .prefetch_related(
                "call_rss", "call_rss__channel", "call_rss__channel__server"
            )
        )

        # ID: Count
        servers_map: dict[int, int] = {}
        servers: dict[int, models.Server] = {}

        for call in calls:
            for rs in utils.deduplicate(call.call_rss, lambda i: i.channel.server_id):
                server_id = rs.channel.server_id
                servers[server_id] = rs.channel.server

                if server_id in servers_map:
                    servers_map[server_id] += 1

                else:
                    servers_map[server_id] = 1

        server_counts = sorted(servers_map.items(), key=lambda i: i[1])

        for pos, (server_id, calls_count) in enumerate(server_counts):
            servers[server_id].leaderboard_calls_count = calls_count
            servers[server_id].leaderboard_position = pos + 1

        await models.Server.bulk_update(
            servers.values(),
            ["leaderboard_calls_count", "leaderboard_position"],
            batch_size=500,
        )

    @commands.Cog.listener()
    async def on_guild_update(
        self, _before: discord.Guild, after: discord.Guild
    ) -> None:
        """Updates a server.

        Args:
            _before (discord.Guild): The server before it being updated.
            after (discord.Guild): The server after it being updated.
        """

        vanity_code = None

        try:
            vanity = await after.vanity_invite()

            if vanity is not None:
                vanity_code = vanity.code

        except errors.DiscordException:
            pass

        await models.Server.update_or_create(
            {"name": after.name, "vanity_code": vanity_code}, id=after.id
        )

    @commands.group(invoke_without_command=True)
    @commands.guild_only()
    @commands.has_guild_permissions(manage_guild=True)
    async def config(
        self, ctx: commands.Context, setting: str = "", value: str | None = None
    ) -> None:
        """
        Set up Payphone for the current server.

        Settings:
        - `prefix`: Set a custom prefix.
        - `enable`: Enable a (or all) channel(s) to run commands in it. Non-enabled channels will pass to be disabled.
        - `disable`: Disables a channel to run commands in it. Non-disabled channels will stay enabled by default.

        You can use `{prefix}help config <setting>` to get help about a specific configuration.
        """

        await ctx.send_help("config")

    @config.command(extras={"help_name": "Prefix"})
    @commands.guild_only()
    @commands.has_guild_permissions(manage_guild=True)
    @transactions.atomic()
    async def prefix(self, ctx: commands.Context, prefix: str = "p.") -> None:
        """
        Set a custom bot prefix.

        Options:
        - `<prefix>`: The new custom prefix.

        Permissions:
        - Manage Server
        """

        try:
            await self.bot.set_prefix(ctx.guild, prefix)
        except ValueError:
            await ctx.send("❌ The prefix must not be longer than 5 characters.")
            return

        await ctx.send(f"✅ The prefix was set to `{prefix}`!")

    @config.command(extras={"help_name": "Enable Channel"})
    @commands.guild_only()
    @commands.has_guild_permissions(manage_channels=True)
    @transactions.atomic()
    async def enable(
        self, ctx: commands.Context, channel: discord.TextChannel | str | None
    ) -> None:
        """
        Enable a channel to run commands in it. Non-enabled channels will pass to be disabled.

        Options:
        - `[<channel>]`: The channel to enable or "all". If no channel is specified, the current channel will be enabled.

        Permissions:
        - Manage Channels
        """

        server_id = ctx.guild.id
        channel = channel or ctx.channel

        if isinstance(channel, str) and channel.lower() == "all":
            await models.Channel.filter(server_id=server_id).update(enabled=None)

            await ctx.send(
                "✅ All channels were enabled! You can now run commands in any channel."
            )

        else:
            row = await dbutils.get_channel(channel)

            row.enabled = True
            await row.save()

            await ctx.send(
                f"✅ {channel.mention} was enabled! Other non-enabled channels were disabled."
            )

    @config.command(extras={"help_name": "Disable Channel"})
    @commands.guild_only()
    @commands.has_guild_permissions(manage_channels=True)
    @transactions.atomic()
    async def disable(
        self, ctx: commands.Context, channel: discord.TextChannel | None
    ) -> None:
        """
        Disables a channel to run commands in it. Non-disabled channels will stay enabled by default.

        Options:
        - `[<channel>]`: The channel to disable. If no channel is specified, the current channel will be disabled.

        Permissions:
        - Manage Channels
        """

        channel = channel or ctx.channel

        row = await dbutils.get_channel(channel)

        row.enabled = False
        await row.save()

        await ctx.send(
            f"✅ {channel.mention} was disabled. Commands cannot be ran in this channel anymore."
        )


def setup(b: bot.Payphone) -> None:
    """Load the cog into the bot.

    Args:
        b (bot.Payphone): The bot instance to load the cog into.
    """

    b.add_cog(Admin(b))