tools.py

"""
The MIT License (MIT)

Copyright (c) 2021-present Pycord Development

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
"""
import collections
import copy
import io
import json
import os
import time
from abc import ABC
from functools import cached_property
from typing import Union

import bftools
import discord
from discord.ext import commands
from dotenv import load_dotenv
from tortoise.models import Model
from tortoise import fields, Tortoise


class Config(collections.UserDict):
    def __init__(self, __dict, parent, **kwargs):
        self.parent = parent
        super().__init__(__dict, **kwargs)

    def __getitem__(self, key):
        return copy.deepcopy(super().__getitem__(key))

    def __setitem__(self, key, item):
        self._on_change()
        super().__setitem__(key, item)

    def __delitem__(self, key):
        self._on_change()
        super().__delitem__(key)

    def _on_change(self):
        self.parent.update_config()


class Cache(Config):
    def _on_change(self):
        self.parent.update_cache()


class Context(discord.ApplicationContext):
    async def respond(self, *args, **kwargs):
        default = {
            'allowed_mentions': discord.AllowedMentions.none()
        }
        default.update(**kwargs)
        await super().respond(*args, **kwargs)

    async def send(self, *args, **kwargs):
        default = {
            'allowed_mentions': discord.AllowedMentions.none()
        }
        default.update(**kwargs)
        await super().send(*args, **kwargs)


class Storage:
    def __init__(self, storage_dir="storage"):
        self._initialized = False
        self.storage_dir = storage_dir
        self.config = None
        self.cache = None
        if not os.path.exists(self.storage_dir):
            os.mkdir(self.storage_dir)
        for filename in ("config", "cache"):
            if not os.path.exists(f"{self.storage_dir}/{filename}.json"):
                open(f"{self.storage_dir}/{filename}.json", "x")
                open(f"{self.storage_dir}/{filename}.json", "w").write("{}")
        self.load_config()
        self.load_cache()
        self._initialized = True

    async def setup_db(self):
        await Tortoise.init(
            db_url=f'sqlite://{self.storage_dir}/main.db',
            modules={'models': ['tools']})
        await Tortoise.generate_schemas()

    def load_config(self):
        with open(f"{self.storage_dir}/config.json", "r") as f:
            data = json.load(f)
            self.config = Config(data, self)
            return self.config

    def load_cache(self):
        with open(f"{self.storage_dir}/cache.json", "r") as f:
            data = json.load(f)
            self.cache = Cache(data, self)
            return self.cache

    def update_config(self):
        try:
            data = self.config.data
            with open(f"{self.storage_dir}/config.json", "w") as f:
                json.dump(data, f, indent=4)
        except AttributeError:
            if self._initialized:
                raise

    def update_cache(self):
        try:
            data = self.cache.data
            with open(f"{self.storage_dir}/cache.json", "w") as f:
                json.dump(data, f, indent=4)
        except AttributeError:
            if self._initialized:
                raise


class Bot(commands.Bot, ABC):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        load_dotenv()
        self.token = os.getenv("BOT_TOKEN")
        self.default_owner = os.getenv('OWNER_ID')
        os.environ["JISHAKU_NO_UNDERSCORE"] = "True"
        os.environ['JISHAKU_RETAIN'] = "True"
        self.owner_id = None
        self.storage = Storage()
        self.owner_ids = self.config.get('owner_ids', [690420846774321221, 556119013298667520])
        self.load_extension('jishaku')
        self.brainfuck = bftools.BrainfuckTools()
        self.hang = False

    @property
    def config(self):
        return self.storage.config

    @property
    def cache(self):
        return self.storage.cache

    def run(self, *args, **kwargs):
        if len(args):
            super().run(*args, **kwargs)
        else:
            super().run(self.token, **kwargs)

    async def get_application_context(
        self, interaction, cls=Context
    ):
        return await super().get_application_context(interaction=interaction, cls=cls)


def escape(text):
    return text.replace("`" * 3, "`​``")


def codeblock(code, lang="py"):
    return f"```{lang}\n{escape(code)}\n```"


def codefile(code, filename=None, ext="py"):
    if filename is None:
        filename = f"code.{ext}"
    data = io.BytesIO(bytes(code, encoding='utf8'))
    return discord.File(data, filename)


async def send_code(ctx, code, lang="py", filename=None, ext=None):
    cb = codeblock(code, lang=lang)
    if len(cb) < 500:  # max is 2000 but using 500 minimizes flood
        return await ctx.respond(cb, allowed_mentions=discord.AllowedMentions.none())
    else:
        if ext is None:
            ext = lang
        cf = codefile(code, filename=filename, ext=ext)
        await ctx.send(file=cf)
        await ctx.respond("Text was too long to put in a codeblock, used file instead")


class Timer:
    def __init__(self):
        self.start_time = time.perf_counter()
        self.end_time = None

    @cached_property
    def message(self):
        return f"Finished in {(self.duration * 1000):.0f}ms"

    @cached_property
    def duration(self):
        return self.end_time - self.start_time

    def finish(self):
        self.end_time = time.perf_counter()
        return self.duration


async def get_prefix(bot, message):
    # TODO: custom prefixes
    return commands.when_mentioned_or(bot.config.get('prefix', ';'))(bot, message)


class Tag(Model):
    name = fields.CharField(null=False, max_length=100)
    guild = fields.IntField(null=False)
    author = fields.IntField(null=False)
    content = fields.CharField(null=False, max_length=2000)
    created = fields.DatetimeField(auto_now_add=True, null=False)
    edited = fields.DatetimeField(auto_now=True, null=False)
    uses = fields.IntField(null=False, default=0)

    def raw_content(self):
        return discord.utils.escape_markdown(self.content)


class Lowercase:
    def __init__(self):
        pass

    async def convert(self, text):
        return text.lower()