miyu-bot/miyu_bot/commands/common/master_asset_manager.py

import hashlib
from functools import lru_cache
from timeit import default_timer
from typing import Callable, Any, Optional, Union

from d4dj_utils.manager.asset_manager import AssetManager
from d4dj_utils.master.master_asset import MasterDict, MasterAsset
from discord.ext import commands

from miyu_bot.commands.common.aliases.event import event_aliases
from miyu_bot.commands.common.fuzzy_matching import FuzzyFilteredMap, romanize

import datetime as dt


class MasterFilterManager:
    def __init__(self, manager: AssetManager):
        self.manager = manager
        self.music = MasterFilter(
            self.manager.music_master,
            naming_function=lambda m: f'{m.name} {m.special_unit_name}',
            filter_function=lambda m: m.is_released,
            fallback_naming_function=lambda m: m.id,
        )
        self.events = MasterFilter(
            self.manager.event_master,
            aliases=event_aliases,
            naming_function=lambda e: e.name,
            filter_function=lambda e: e.start_datetime < dt.datetime.now(dt.timezone.utc) + dt.timedelta(hours=12),
        )
        self.cards = MasterFilter(
            self.manager.card_master,
            naming_function=lambda c: c.name,
            filter_function=lambda c: c.is_released,
        )


class MasterFilter:
    def __init__(self, masters: MasterDict,
                 naming_function: Callable[[Any], str],
                 aliases: Optional[dict] = None,
                 filter_function=lambda _: True,
                 fallback_naming_function: Optional[Callable[[Any], str]] = None):
        self.masters = masters
        self.default_filter = FuzzyFilteredMap(filter_function)
        self.unrestricted_filter = FuzzyFilteredMap()
        for master in masters.values():
            name = romanize(naming_function(master))
            if fallback_naming_function and self.default_filter.has_exact_unprocessed(name):
                name = romanize(fallback_naming_function(master))
                if self.default_filter.has_exact_unprocessed(name):
                    continue
            elif self.default_filter.has_exact_unprocessed(name):
                continue
            self.default_filter.set_unprocessed(name, master)
            self.unrestricted_filter.set_unprocessed(name, master)
        if aliases:
            for alias, mid in aliases.items():
                master = masters[mid]
                alias = romanize(alias)
                self.default_filter.set_unprocessed(alias, master)
                self.unrestricted_filter.set_unprocessed(alias, master)

    def get(self, name_or_id: Union[str, int], ctx: Optional[commands.Context]):
        if ctx and ctx.channel.id in no_filter_channels:
            try:
                return self.masters[int(name_or_id)]
            except (KeyError, ValueError):
                if isinstance(name_or_id, int):
                    return None
                return self.unrestricted_filter[name_or_id]
        else:
            try:
                master = self.masters[int(name_or_id)]
                if master not in self.default_filter.values():
                    master = self.default_filter[name_or_id]
                return master
            except (KeyError, ValueError):
                if isinstance(name_or_id, int):
                    return None
                return self.default_filter[name_or_id]

    def get_sorted(self, name: str, ctx: commands.Context):
        if name:
            if ctx.channel.id in no_filter_channels:
                return self.unrestricted_filter.get_sorted(name)
            else:
                return self.default_filter.get_sorted(name)
        else:
            if ctx.channel.id in no_filter_channels:
                return list(self.unrestricted_filter.values())
            else:
                return list(self.default_filter.values())

    def values(self, ctx: commands.Context):
        if ctx.channel.id in no_filter_channels:
            return self.unrestricted_filter.values()
        else:
            return self.default_filter.values()


def hash_master(master: MasterAsset):
    return hashlib.md5(master.extended_description().encode('utf-8')).hexdigest()


no_filter_channels = {790033228600705048, 790033272376918027, 795640603114864640}