Red-DiscordBot/redbot/core/settings_caches.py

from __future__ import annotations

from typing import Dict, List, Optional, Union, Set, Iterable, Tuple, overload
import asyncio
from argparse import Namespace
from collections import defaultdict

import discord

from .config import Config
from .utils import AsyncIter


class PrefixManager:
    def __init__(self, config: Config, cli_flags: Namespace):
        self._config: Config = config
        self._global_prefix_overide: Optional[List[str]] = (
            sorted(cli_flags.prefix, reverse=True) or None
        )
        self._cached: Dict[Optional[int], List[str]] = {}

    async def get_prefixes(self, guild: Optional[discord.Guild] = None) -> List[str]:
        ret: List[str]

        gid: Optional[int] = guild.id if guild else None

        if gid in self._cached:
            ret = self._cached[gid].copy()
        else:
            if gid is not None:
                ret = await self._config.guild_from_id(gid).prefix()
                if not ret:
                    ret = await self.get_prefixes(None)
            else:
                ret = self._global_prefix_overide or (await self._config.prefix())

            self._cached[gid] = ret.copy()

        return ret

    async def set_prefixes(
        self, guild: Optional[discord.Guild] = None, prefixes: Optional[List[str]] = None
    ):
        gid: Optional[int] = guild.id if guild else None
        prefixes = prefixes or []
        if not isinstance(prefixes, list) and not all(isinstance(pfx, str) for pfx in prefixes):
            raise TypeError("Prefixes must be a list of strings")
        prefixes = sorted(prefixes, reverse=True)
        if gid is None:
            if not prefixes:
                raise ValueError("You must have at least one prefix.")
            self._cached.clear()
            await self._config.prefix.set(prefixes)
        else:
            self._cached.pop(gid, None)
            await self._config.guild_from_id(gid).prefix.set(prefixes)


class I18nManager:
    def __init__(self, config: Config):
        self._config: Config = config
        self._guild_locale: Dict[Union[int, None], Union[str, None]] = {}
        self._guild_regional_format: Dict[Union[int, None], Union[str, None]] = {}

    async def get_locale(self, guild: Union[discord.Guild, None]) -> str:
        """Get the guild locale from the cache"""
        # Ensure global locale is in the cache
        if None not in self._guild_locale:
            global_locale = await self._config.locale()
            self._guild_locale[None] = global_locale

        if guild is None:  # Not a guild so cannot support guild locale
            # Return the bot's globally set locale if its None on a guild scope.
            return self._guild_locale[None]
        elif guild.id in self._guild_locale:  # Cached guild
            if self._guild_locale[guild.id] is None:
                return self._guild_locale[None]
            else:
                return self._guild_locale[guild.id]
        else:  # Uncached guild
            out = await self._config.guild(guild).locale()  # No locale set
            if out is None:
                self._guild_locale[guild.id] = None
                return self._guild_locale[None]
            else:
                self._guild_locale[guild.id] = out
                return out

    @overload
    async def set_locale(self, guild: None, locale: str):
        ...

    @overload
    async def set_locale(self, guild: discord.Guild, locale: Union[str, None]):
        ...

    async def set_locale(
        self, guild: Union[discord.Guild, None], locale: Union[str, None]
    ) -> None:
        """Set the locale in the config and cache"""
        if guild is None:
            if locale is None:
                # this method should never be called like this
                raise ValueError("Global locale can't be None!")
            self._guild_locale[None] = locale
            await self._config.locale.set(locale)
            return
        self._guild_locale[guild.id] = locale
        await self._config.guild(guild).locale.set(locale)

    async def get_regional_format(self, guild: Union[discord.Guild, None]) -> Optional[str]:
        """Get the regional format from the cache"""
        # Ensure global locale is in the cache
        if None not in self._guild_regional_format:
            global_regional_format = await self._config.regional_format()
            self._guild_regional_format[None] = global_regional_format

        if guild is None:  # Not a guild so cannot support guild locale
            return self._guild_regional_format[None]
        elif guild.id in self._guild_regional_format:  # Cached guild
            if self._guild_regional_format[guild.id] is None:
                return self._guild_regional_format[None]
            else:
                return self._guild_regional_format[guild.id]
        else:  # Uncached guild
            out = await self._config.guild(guild).regional_format()  # No locale set
            if out is None:
                self._guild_regional_format[guild.id] = None
                return self._guild_regional_format[None]
            else:  # Not cached, got a custom regional format.
                self._guild_regional_format[guild.id] = out
                return out

    async def set_regional_format(
        self, guild: Union[discord.Guild, None], regional_format: Union[str, None]
    ) -> None:
        """Set the regional format in the config and cache"""
        if guild is None:
            self._guild_regional_format[None] = regional_format
            await self._config.regional_format.set(regional_format)
            return
        self._guild_regional_format[guild.id] = regional_format
        await self._config.guild(guild).regional_format.set(regional_format)


class IgnoreManager:
    def __init__(self, config: Config):
        self._config: Config = config
        self._cached_channels: Dict[int, bool] = {}
        self._cached_guilds: Dict[int, bool] = {}

    async def get_ignored_channel(
        self, channel: Union[discord.TextChannel, discord.Thread], check_category: bool = True
    ) -> bool:
        ret: bool

        cid: int = channel.id
        cat_id: Optional[int] = (
            channel.category.id if check_category and channel.category else None
        )
        if cid in self._cached_channels:
            chan_ret = self._cached_channels[cid]
        else:
            chan_ret = await self._config.channel_from_id(cid).ignored()
            self._cached_channels[cid] = chan_ret
        if cat_id and cat_id in self._cached_channels:
            cat_ret = self._cached_channels[cat_id]
        else:
            if cat_id:
                cat_ret = await self._config.channel_from_id(cat_id).ignored()
                self._cached_channels[cat_id] = cat_ret
            else:
                cat_ret = False
        ret = chan_ret or cat_ret

        return ret

    async def set_ignored_channel(
        self,
        channel: Union[discord.TextChannel, discord.Thread, discord.CategoryChannel],
        set_to: bool,
    ):
        cid: int = channel.id
        self._cached_channels[cid] = set_to
        if set_to:
            await self._config.channel_from_id(cid).ignored.set(set_to)
        else:
            await self._config.channel_from_id(cid).ignored.clear()

    async def get_ignored_guild(self, guild: discord.Guild) -> bool:
        ret: bool

        gid: int = guild.id

        if gid in self._cached_guilds:
            ret = self._cached_guilds[gid]
        else:
            ret = await self._config.guild_from_id(gid).ignored()
            self._cached_guilds[gid] = ret

        return ret

    async def set_ignored_guild(self, guild: discord.Guild, set_to: bool):
        gid: int = guild.id
        self._cached_guilds[gid] = set_to
        if set_to:
            await self._config.guild_from_id(gid).ignored.set(set_to)
        else:
            await self._config.guild_from_id(gid).ignored.clear()


class WhitelistBlacklistManager:
    def __init__(self, config: Config):
        self._config: Config = config
        self._cached_whitelist: Dict[Optional[int], Set[int]] = {}
        self._cached_blacklist: Dict[Optional[int], Set[int]] = {}
        # because of discord deletion
        # we now have sync and async access that may need to happen at the
        # same time.
        # blame discord for this.
        self._access_lock = asyncio.Lock()

    async def discord_deleted_user(self, user_id: int):
        async with self._access_lock:
            async for guild_id_or_none, ids in AsyncIter(
                self._cached_whitelist.items(), steps=100
            ):
                ids.discard(user_id)

            async for guild_id_or_none, ids in AsyncIter(
                self._cached_blacklist.items(), steps=100
            ):
                ids.discard(user_id)

            for grp in (self._config.whitelist, self._config.blacklist):
                async with grp() as ul:
                    try:
                        ul.remove(user_id)
                    except ValueError:
                        pass

            # don't use this in extensions, it's optimized and controlled for here,
            # but can't be safe in 3rd party use

            async with self._config._get_base_group("GUILD").all() as abuse:
                for guild_str, guild_data in abuse.items():
                    for l_name in ("whitelist", "blacklist"):
                        try:
                            guild_data[l_name].remove(user_id)
                        except (ValueError, KeyError):
                            pass  # this is raw access not filled with defaults

    async def get_whitelist(self, guild: Optional[discord.Guild] = None) -> Set[int]:
        async with self._access_lock:
            ret: Set[int]
            gid: Optional[int] = guild.id if guild else None
            if gid in self._cached_whitelist:
                ret = self._cached_whitelist[gid].copy()
            else:
                if gid is not None:
                    ret = set(await self._config.guild_from_id(gid).whitelist())
                else:
                    ret = set(await self._config.whitelist())

                self._cached_whitelist[gid] = ret.copy()

            return ret

    async def add_to_whitelist(self, guild: Optional[discord.Guild], role_or_user: Iterable[int]):
        async with self._access_lock:
            gid: Optional[int] = guild.id if guild else None
            role_or_user = role_or_user or []
            if not all(isinstance(r_or_u, int) for r_or_u in role_or_user):
                raise TypeError("`role_or_user` must be an iterable of `int`s.")

            if gid is None:
                if gid not in self._cached_whitelist:
                    self._cached_whitelist[gid] = set(await self._config.whitelist())
                self._cached_whitelist[gid].update(role_or_user)
                await self._config.whitelist.set(list(self._cached_whitelist[gid]))

            else:
                if gid not in self._cached_whitelist:
                    self._cached_whitelist[gid] = set(
                        await self._config.guild_from_id(gid).whitelist()
                    )
                self._cached_whitelist[gid].update(role_or_user)
                await self._config.guild_from_id(gid).whitelist.set(
                    list(self._cached_whitelist[gid])
                )

    async def clear_whitelist(self, guild: Optional[discord.Guild] = None):
        async with self._access_lock:
            gid: Optional[int] = guild.id if guild else None
            self._cached_whitelist[gid] = set()
            if gid is None:
                await self._config.whitelist.clear()
            else:
                await self._config.guild_from_id(gid).whitelist.clear()

    async def remove_from_whitelist(
        self, guild: Optional[discord.Guild], role_or_user: Iterable[int]
    ):
        async with self._access_lock:
            gid: Optional[int] = guild.id if guild else None
            role_or_user = role_or_user or []
            if not all(isinstance(r_or_u, int) for r_or_u in role_or_user):
                raise TypeError("`role_or_user` must be an iterable of `int`s.")

            if gid is None:
                if gid not in self._cached_whitelist:
                    self._cached_whitelist[gid] = set(await self._config.whitelist())
                self._cached_whitelist[gid].difference_update(role_or_user)
                await self._config.whitelist.set(list(self._cached_whitelist[gid]))

            else:
                if gid not in self._cached_whitelist:
                    self._cached_whitelist[gid] = set(
                        await self._config.guild_from_id(gid).whitelist()
                    )
                self._cached_whitelist[gid].difference_update(role_or_user)
                await self._config.guild_from_id(gid).whitelist.set(
                    list(self._cached_whitelist[gid])
                )

    async def get_blacklist(self, guild: Optional[discord.Guild] = None) -> Set[int]:
        async with self._access_lock:
            ret: Set[int]
            gid: Optional[int] = guild.id if guild else None
            if gid in self._cached_blacklist:
                ret = self._cached_blacklist[gid].copy()
            else:
                if gid is not None:
                    ret = set(await self._config.guild_from_id(gid).blacklist())
                else:
                    ret = set(await self._config.blacklist())

                self._cached_blacklist[gid] = ret.copy()

            return ret

    async def add_to_blacklist(self, guild: Optional[discord.Guild], role_or_user: Iterable[int]):
        async with self._access_lock:
            gid: Optional[int] = guild.id if guild else None
            role_or_user = role_or_user or []
            if not all(isinstance(r_or_u, int) for r_or_u in role_or_user):
                raise TypeError("`role_or_user` must be an iterable of `int`s.")
            if gid is None:
                if gid not in self._cached_blacklist:
                    self._cached_blacklist[gid] = set(await self._config.blacklist())
                self._cached_blacklist[gid].update(role_or_user)
                await self._config.blacklist.set(list(self._cached_blacklist[gid]))
            else:
                if gid not in self._cached_blacklist:
                    self._cached_blacklist[gid] = set(
                        await self._config.guild_from_id(gid).blacklist()
                    )
                self._cached_blacklist[gid].update(role_or_user)
                await self._config.guild_from_id(gid).blacklist.set(
                    list(self._cached_blacklist[gid])
                )

    async def clear_blacklist(self, guild: Optional[discord.Guild] = None):
        async with self._access_lock:
            gid: Optional[int] = guild.id if guild else None
            self._cached_blacklist[gid] = set()
            if gid is None:
                await self._config.blacklist.clear()
            else:
                await self._config.guild_from_id(gid).blacklist.clear()

    async def remove_from_blacklist(
        self, guild: Optional[discord.Guild], role_or_user: Iterable[int]
    ):
        async with self._access_lock:
            gid: Optional[int] = guild.id if guild else None
            role_or_user = role_or_user or []
            if not all(isinstance(r_or_u, int) for r_or_u in role_or_user):
                raise TypeError("`role_or_user` must be an iterable of `int`s.")
            if gid is None:
                if gid not in self._cached_blacklist:
                    self._cached_blacklist[gid] = set(await self._config.blacklist())
                self._cached_blacklist[gid].difference_update(role_or_user)
                await self._config.blacklist.set(list(self._cached_blacklist[gid]))
            else:
                if gid not in self._cached_blacklist:
                    self._cached_blacklist[gid] = set(
                        await self._config.guild_from_id(gid).blacklist()
                    )
                self._cached_blacklist[gid].difference_update(role_or_user)
                await self._config.guild_from_id(gid).blacklist.set(
                    list(self._cached_blacklist[gid])
                )


class DisabledCogCache:
    def __init__(self, config: Config):
        self._config = config
        self._disable_map: Dict[str, Dict[int, bool]] = defaultdict(dict)

    async def cog_disabled_in_guild(self, cog_name: str, guild_id: int) -> bool:
        """
        Check if a cog is disabled in a guild

        Parameters
        ----------
        cog_name: str
            This should be the cog's qualified name, not necessarily the classname
        guild_id: int

        Returns
        -------
        bool
        """

        if guild_id in self._disable_map[cog_name]:
            return self._disable_map[cog_name][guild_id]

        gset = await self._config.custom("COG_DISABLE_SETTINGS", cog_name, guild_id).disabled()
        if gset is None:
            gset = await self._config.custom("COG_DISABLE_SETTINGS", cog_name, 0).disabled()
            if gset is None:
                gset = False

        self._disable_map[cog_name][guild_id] = gset
        return gset

    async def default_disable(self, cog_name: str):
        """
        Sets the default for a cog as disabled.

        Parameters
        ----------
        cog_name: str
            This should be the cog's qualified name, not necessarily the classname
        """
        await self._config.custom("COG_DISABLE_SETTINGS", cog_name, 0).disabled.set(True)
        self._disable_map.pop(cog_name, None)

    async def default_enable(self, cog_name: str):
        """
        Sets the default for a cog as enabled.

        Parameters
        ----------
        cog_name: str
            This should be the cog's qualified name, not necessarily the classname
        """
        await self._config.custom("COG_DISABLE_SETTINGS", cog_name, 0).disabled.clear()
        self._disable_map.pop(cog_name, None)

    async def disable_cog_in_guild(self, cog_name: str, guild_id: int) -> bool:
        """
        Disable a cog in a guild.

        Parameters
        ----------
        cog_name: str
            This should be the cog's qualified name, not necessarily the classname
        guild_id: int

        Returns
        -------
        bool
            Whether or not any change was made.
            This may be useful for settings commands.
        """

        if await self.cog_disabled_in_guild(cog_name, guild_id):
            return False

        self._disable_map[cog_name][guild_id] = True
        await self._config.custom("COG_DISABLE_SETTINGS", cog_name, guild_id).disabled.set(True)
        return True

    async def enable_cog_in_guild(self, cog_name: str, guild_id: int) -> bool:
        """
        Enable a cog in a guild.

        Parameters
        ----------
        cog_name: str
            This should be the cog's qualified name, not necessarily the classname
        guild_id: int

        Returns
        -------
        bool
            Whether or not any change was made.
            This may be useful for settings commands.
        """

        if not await self.cog_disabled_in_guild(cog_name, guild_id):
            return False

        self._disable_map[cog_name][guild_id] = False
        await self._config.custom("COG_DISABLE_SETTINGS", cog_name, guild_id).disabled.set(False)
        return True