From 87d6babcc86245c8c147daad6f20d0bd512b6592 Mon Sep 17 00:00:00 2001
From: jack1142 <6032823+jack1142@users.noreply.github.com>
Date: Tue, 20 Apr 2021 14:00:41 +0200
Subject: [PATCH] [Maintenance] Add `--end-after-restart` option

---
 maintenance/__init__.py    |  1 +
 maintenance/classes.py     |  4 +++-
 maintenance/converters.py  |  3 +++
 maintenance/maintenance.py | 35 +++++++++++++++++++++++++----------
 4 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/maintenance/__init__.py b/maintenance/__init__.py
index 48116608..4f8010cf 100644
--- a/maintenance/__init__.py
+++ b/maintenance/__init__.py
@@ -10,3 +10,4 @@
 async def setup(bot):
     cog = Maintenance(bot)
     await bot.add_cog(cog)
+    await cog.initialize()
diff --git a/maintenance/classes.py b/maintenance/classes.py
index 0ca2a5de..32ceb7de 100644
--- a/maintenance/classes.py
+++ b/maintenance/classes.py
@@ -6,7 +6,9 @@ def __init__(self, start=0, end=0, after=True, whitelist=None):
         if not whitelist:
             whitelist = []
         self.start = start + time.time()
-        if end:
+        if end is True:
+            self.end = True
+        elif end:
             if after:
                 self.end = end + self.start
             else:
diff --git a/maintenance/converters.py b/maintenance/converters.py
index 62642ee8..f1e095b6 100644
--- a/maintenance/converters.py
+++ b/maintenance/converters.py
@@ -23,6 +23,7 @@ async def convert(self, ctx, argument):
         _end = parser.add_mutually_exclusive_group()
         _end.add_argument("--end-after", nargs="*", dest="end", default=[])
         _end.add_argument("--end-in", nargs="*", dest="endin", default=[])
+        _end.add_argument("--end-after-restart", dest="end_after_restart", action="store_true")
         try:
             vals = vars(parser.parse_args(argument.split(" ")))
         except Exception as exc:
@@ -35,6 +36,8 @@ async def convert(self, ctx, argument):
         if end_seconds == None:
             end_seconds = convert_time(vals.get("endin", None))
             after = False
+        if vals.get("end_after_restart", False):
+            end_seconds = True
         if start_seconds:
             if end_seconds:
                 scheduled = ScheduledMaintenance(
diff --git a/maintenance/maintenance.py b/maintenance/maintenance.py
index c1594360..d11108e7 100644
--- a/maintenance/maintenance.py
+++ b/maintenance/maintenance.py
@@ -58,10 +58,21 @@ def __init__(self, bot):
             "delete": 3,
             "scheduledmaintenance": [],
         }
-        # When on maintenance, on will be set to [True, second of when it is off maintenance, list of people who can bypass the maintenance]
+        # When on maintenance, on will be set to [True, second of when it is off maintenance or True if it's planned to end after bot restart, list of people who can bypass the maintenance]
         self.conf.register_global(**default_global)
         self.bot.add_check(self.this_check)
-        self.task = self.bot.loop.create_task(self.bg_loop())
+        self.task = None
+
+    async def cog_load(self):
+        # disable maintenance scheduled to end after restart,
+        # if the cog was loaded during bot startup
+        if not self.bot.is_ready():
+            on = await self.conf.on()
+            if on[0] and on[1] is True:
+                await self.conf.on.set([False, 0, []])
+
+    async def initialize(self):
+        self.task = asyncio.create_task(self.bg_loop())
 
     @listener()
     async def on_command_error(self, ctx, error):
@@ -73,11 +84,9 @@ async def on_command_error(self, ctx, error):
                 await ctx.send(error)
 
     def cog_unload(self):
-        self.__unload()
-
-    def __unload(self):
         self.bot.remove_check(self.this_check)
-        self.task.cancel()
+        if self.task is not None:
+            self.task.cancel()
 
     async def red_delete_data_for_user(
         self,
@@ -121,7 +130,10 @@ async def this_check(self, ctx):
         on = await self.conf.on()
         if not on[0]:
             return True
-        if on[1] <= time.time() and on[1] != 0:
+        if on[1] is True:
+            # maintenance will be disabled after restart
+            pass
+        elif on[1] != 0 and on[1] <= time.time():
             setting = [False, 0, []]
             await self.conf.on.set(setting)
             return True
@@ -145,6 +157,7 @@ async def _on(self, ctx, *, args: Margs = None):
             --start-in: Makes the maintenace start in that long.
             --end-in: Schedules the maintenance to end in that long from the current second.
             --end-after: Schedules the maintenance to end in that long after the maitenance has started.
+            --end-after-restart: Schedules the maintenance to end after bot restart.
             --whitelist: Provide user IDs after this to whitelist people from the maintenance.
 
         Examples:
@@ -165,8 +178,8 @@ async def _on(self, ctx, *, args: Margs = None):
         else:
             num = 0
             whitelist = []
-            start = time.time()
-        if start == time.time():
+            start = None
+        if start is None:
             setting = [True, num, whitelist]
             await self.conf.on.set(setting)
         else:
@@ -195,7 +208,9 @@ async def settings(self, ctx):
                     sending += "    • " + starting
                 sending += "```"
             return await ctx.send(sending)
-        if on[1] != 0:
+        if on[1] is True:
+            done = "after bot restart"
+        elif on[1] != 0:
             done = str(datetime.fromtimestamp(on[1]).strftime("%A, %B %d, %Y %I:%M:%S"))
             done = "on " + done
         else: