Fox-V3/fifo/fifo.py

from datetime import datetime, timedelta
from typing import Dict, Union

import discord
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.base import BaseTrigger
from apscheduler.triggers.combining import OrTrigger
from apscheduler.triggers.date import DateTrigger
from apscheduler.triggers.interval import IntervalTrigger
from dateutil import parser
from redbot.core import Config, checks, commands
from redbot.core.bot import Red
from redbot.core.commands import DictConverter, TimedeltaConverter, parse_timedelta

from .datetimeconverter import DatetimeConverter
from .redconfigjobstore import RedConfigJobStore


def get_trigger(data):
    if data["type"] == "interval":
        parsed_time = data["time_data"]
        return IntervalTrigger(days=parsed_time.days, seconds=parsed_time.seconds)

    if data["type"] == "date":
        return DateTrigger(data["time_data"])

    if data["type"] == "cron":
        return None  # TODO: Cron parsing

    return False


def parse_triggers(data: Union[Dict, None]):
    if data is None or not data.get("triggers", False):  # No triggers
        return None

    if len(data["triggers"]) > 1:  # Multiple triggers
        return OrTrigger(get_trigger(t_data) for t_data in data["triggers"])

    return get_trigger(data[0])


class FakeMessage:
    _state = None


# class FakeMessage(discord.Message):
#     def __init__(self):
#         super().__init__(state=None, channel=None, data=None)


class Task:
    default_task_data = {"triggers": [], "command_str": ""}

    default_trigger = {
        "type": "",
        "time_data": None,  # Used for Interval and Date Triggers
    }

    def __init__(self, name: str, guild_id, config: Config, author_id=None, bot: Red = None):
        self.name = name
        self.guild_id = guild_id
        self.config = config
        self.bot = bot
        self.author_id = author_id
        self.data = None

    async def _encode_time_data(self):
        if not self.data or not self.data.get("triggers", None):
            return None

        triggers = []
        for t in self.data["triggers"]:
            if t["type"] == "interval":  # Convert into timedelta
                td: timedelta = t["time_data"]

                triggers.append({"type": t["type"], "time_data":  {"days": td.days, "seconds": td.seconds} })

            if t["type"] == "date":  # Convert into datetime
                dt: datetime = t["time_data"]
                triggers.append({"type": t["type"], "time_data":  {
                    "year": dt.year,
                    "month": dt.month,
                    "day": dt.day,
                    "hour": dt.hour,
                    "minute": dt.minute,
                    "second": dt.second,
                }})

            if t["type"] == "cron":
                raise NotImplemented
            raise NotImplemented

        return triggers

    async def _decode_time_data(self):
        if not self.data or not self.data.get("triggers", None):
            return

        for t in self.data["triggers"]:
            if t["type"] == "interval":  # Convert into timedelta
                t["time_data"] = timedelta(**t["time_data"])

            if t["type"] == "date":  # Convert into datetime
                t["time_data"] = datetime(**t["time_data"])

            if t["type"] == "cron":
                raise NotImplemented
            raise NotImplemented

    # async def load_from_data(self, data: Dict):
    #     self.data = data.copy()

    async def load_from_config(self):
        data = await self.config.guild_from_id(self.guild_id).tasks.get_raw(
            self.name, default=None
        )

        if not data:
            return

        self.author_id = data["author_id"]
        self.guild_id = data["guild_id"]

        self.data = data["data"]

        await self._decode_time_data()
        return self.data

    async def get_trigger(self) -> Union[BaseTrigger, None]:
        if not self.data:
            await self.load_from_config()

        return parse_triggers(self.data)

    # async def set_job_id(self, job_id):
    #     if self.data is None:
    #         await self.load_from_config()
    #
    #     self.data["job_id"] = job_id

    async def save_all(self):
        """To be used when creating an new task"""

        data_to_save = self.default_task_data.copy()
        if self.data:
            data_to_save["command_str"] = self.data.get("command_str", "")
            data_to_save["triggers"] = await self._encode_time_data()

        to_save = {
            "guild_id": self.guild_id,
            "author_id": self.author_id,
            "data": data_to_save,
        }
        await self.config.guild_from_id(self.guild_id).tasks.set_raw(self.name, value=to_save)

    async def save_data(self):
        """To be used when updating triggers"""
        if not self.data:
            return
        await self.config.guild_from_id(self.guild_id).tasks.set_raw(
            self.name, "data", value=await self._encode_time_data()
        )

    async def execute(self):
        if not self.data or self.data["command_str"]:
            return False
        message = FakeMessage()
        message.guild = self.bot.get_guild(self.guild_id)  # used for get_prefix
        message.author = message.guild.get_member(self.author_id)
        message.content = await self.bot.get_prefix(message) + self.data["command_str"]

        if not message.guild or not message.author or not message.content:
            return False

        new_ctx: commands.Context = await self.bot.get_context(message)
        if not new_ctx.valid:
            return False

        await self.bot.invoke(new_ctx)
        return True

    async def set_bot(self, bot: Red):
        self.bot = bot

    async def set_author(self, author: Union[discord.User, str]):
        self.author_id = getattr(author, "id", None) or author

    async def set_commmand_str(self, command_str):
        if not self.data:
            self.data = self.default_task_data.copy()
        self.data["command_str"] = command_str
        return True

    async def add_trigger(self, param, parsed_time: Union[timedelta, datetime]):
        trigger_data = {"type": param, "time_data": parsed_time}
        if not get_trigger(trigger_data):
            return False

        if not self.data:
            self.data = self.default_task_data.copy()

        self.data["triggers"].append(trigger_data)
        return True


class FIFO(commands.Cog):
    """
    Simple Scheduling Cog

    Named after the simplest scheduling algorithm: First In First Out
    """

    def __init__(self, bot: Red):
        super().__init__()
        self.bot = bot
        self.config = Config.get_conf(self, identifier=70737079, force_registration=True)

        default_global = {"jobs_index": {}, "jobs": []}
        default_guild = {"tasks": {}}

        self.config.register_global(**default_global)
        self.config.register_guild(**default_guild)

        jobstores = {"default": RedConfigJobStore(self.config, self.bot)}

        job_defaults = {"coalesce": False, "max_instances": 1}

        # executors = {"default": AsyncIOExecutor()}

        # Default executor is already AsyncIOExecutor
        self.scheduler = AsyncIOScheduler(jobstores=jobstores, job_defaults=job_defaults)

        self.scheduler.start()

    async def red_delete_data_for_user(self, **kwargs):
        """Nothing to delete"""
        return

    def _assemble_job_id(self, task_name, guild_id):
        return task_name + "_" + guild_id

    async def _check_parsable_command(self, ctx: commands.Context, command_to_parse: str):
        message = FakeMessage()
        message.content = ctx.prefix + command_to_parse
        message.author = ctx.author
        message.guild = ctx.guild

        new_ctx: commands.Context = await self.bot.get_context(message)

        return new_ctx.valid

    async def _get_job(self, task_name, guild_id):
        return self.scheduler.get_job(self._assemble_job_id(task_name, guild_id))

    async def _add_job(self, task):
        return self.scheduler.add_job(
            task.execute,
            id=self._assemble_job_id(task.name, task.guild_id),
            trigger=await task.get_trigger(),
        )

    @checks.is_owner()
    @commands.command()
    async def fifoclear(self, ctx: commands.Context):
        """Debug command to clear fifo config"""
        await self.config.guild(ctx.guild).tasks.clear()
        await ctx.tick()

    @checks.is_owner()  # Will be reduced when I figure out permissions later
    @commands.group()
    async def fifo(self, ctx: commands.Context):
        """
        Base command for handling scheduling of tasks
        """
        if ctx.invoked_subcommand is None:
            pass

    @fifo.command(name="list")
    async def fifo_list(self, ctx: commands.Context, all_guilds: bool = False):
        """
        Lists all current tasks and their triggers.

        Do `[p]fifo list True` to see tasks from all guilds
        """
        if all_guilds:
            pass
        else:
            pass  # TODO: parse and display tasks

    @fifo.command(name="add")
    async def fifo_add(self, ctx: commands.Context, task_name: str, *, command_to_execute: str):
        """
        Add a new task to this guild's task list
        """
        if (await self.config.guild(ctx.guild).tasks.get_raw(task_name, default=None)) is not None:
            await ctx.maybe_send_embed(f"Task already exists with {task_name=}")
            return

        if not await self._check_parsable_command(ctx, command_to_execute):
            await ctx.maybe_send_embed("Failed to parse command. Make sure to include the prefix")
            return

        task = Task(task_name, ctx.guild.id, self.config, ctx.author.id)
        await task.set_commmand_str(command_to_execute)
        await task.save_all()
        await ctx.tick()

    @fifo.command(name="delete")
    async def fifo_delete(self, ctx: commands.Context, task_name: str):
        """
        Deletes a task from this guild's task list
        """
        pass

    @fifo.group(name="trigger")
    async def fifo_trigger(self, ctx: commands.Context):
        """
        Add a new trigger for a task from the current guild.
        """
        if ctx.invoked_subcommand is None:
            pass

    @fifo_trigger.command(name="interval")
    async def fifo_trigger_interval(
        self, ctx: commands.Context, task_name: str, interval_str: TimedeltaConverter
    ):
        """
        Add an interval trigger to the specified task
        """

        task = Task(task_name, ctx.guild.id, self.config)
        await task.load_from_config()

        if task.data is None:
            await ctx.maybe_send_embed(
                f"Task by the name of {task_name} is not found in this guild"
            )
            return

        result = await task.add_trigger("interval", interval_str)
        if not result:
            await ctx.maybe_send_embed(
                "Failed to add an interval trigger to this task, see console for logs"
            )
            return
        await task.save_data()
        await ctx.tick()

    @fifo_trigger.command(name="date")
    async def fifo_trigger_date(
        self, ctx: commands.Context, task_name: str, datetime_str: DatetimeConverter
    ):
        """
        Add a "run once" datetime trigger to the specified task
        """

        task = Task(task_name, ctx.guild.id, self.config)
        await task.load_from_config()

        if task.data is None:
            await ctx.maybe_send_embed(
                f"Task by the name of {task_name} is not found in this guild"
            )
            return

        result = await task.add_trigger("date", datetime_str)
        if not result:
            await ctx.maybe_send_embed(
                "Failed to add a date trigger to this task, see console for logs"
            )
            return

        await task.save_data()
        await ctx.tick()

    @fifo_trigger.command(name="cron")
    async def fifo_trigger_cron(
        self, ctx: commands.Context, task_name: str, cron_settings: DictConverter
    ):
        """
        Add a "time of day" trigger to the specified task
        """
        await ctx.maybe_send_embed("Not yet implemented")

    # async def load_tasks(self):
    #     """
    #     Run once on cog load.
    #     """
    #     all_guilds = await self.config.all_guilds()
    #     async for guild_id, guild_data in AsyncIter(all_guilds["tasks"].items(), steps=100):
    #         for task_name, task_data in guild_data["tasks"].items():
    #             task = Task(task_name, guild_id, self.config)
    #             await task.load_from_data(task_data)
    #
    #             job = self.scheduler.add_job(
    #                 task.execute, id=task_name + "_" + guild_id, trigger=await task.get_trigger(),
    #             )
    #
    #     self.scheduler.start()

    # async def parent_loop(self):
    #     await asyncio.sleep(60)
    #     asyncio.create_task(self.process_tasks(datetime.datetime.utcnow()))
    #
    # async def process_tasks(self, now: datetime.datetime):
    #     for task in self.tasks:
    #         pass