From fe00f1d635acbcf3199f6c6d465e897625add446 Mon Sep 17 00:00:00 2001
From: Anon <yanderefedi@proton.me>
Date: Sun, 12 Mar 2023 17:39:02 -0700
Subject: [PATCH] Updated to make project modular

---
 src/main.py | 154 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 149 insertions(+), 5 deletions(-)

diff --git a/src/main.py b/src/main.py
index a9703e7..1a84495 100755
--- a/src/main.py
+++ b/src/main.py
@@ -16,10 +16,154 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+import os
 import sys
 import argparse
 import signal
-import yandere_bot
+import random
+import datetime
+import subprocess
+import FediBot
+from mastodon import MastodonAPIError
+
+# A class that contains all of the rendering information for a single post
+class PickAndRenderFrame:
+	# From Config File
+	profile_name = ""
+	path = ""
+	frames = 0
+	skip = tuple()
+	nsfw = tuple()
+	output_name = ""
+	message = ""
+	message_nsfw = ""
+	render_script = ""
+
+	# After pick
+	picked_frame = None
+	output_name_tr = ""
+
+	def __init__(self, picked, datetime_format):
+		dt_now = datetime.datetime.now()
+		dt_now_str = datetime.datetime.strftime(dt_now, datetime_format)
+		self.profile_name =		picked["profile_name"]
+		self.path =				picked["path"]
+		self.frames =			picked["frames"]
+		self.skip =				picked["skip"]
+		self.nsfw =				picked["nsfw"]
+		self.output_name =		picked["output_name"]
+		self.output_name_tr =	self.output_name
+		self.message =			picked["message"]
+		self.message_nsfw =		picked["message_nsfw"]
+		self.render_script =	picked["render_script"]
+
+		self.picked_frame = self._pick_frame()
+
+		# Shell-like substitutions
+		translations = {
+				"profile_name": self.profile_name,
+				"frame": str(self.picked_frame),
+				"datetime": dt_now_str
+			}
+		
+		self.output_name_tr = self._translate_basename(translations)
+		self._render_frame()
+	
+	def __del__(self):
+		render_file = self.output_name_tr
+		if render_file and os.path.isfile(render_file):
+			os.remove(render_file)
+
+	def _pick_frame(self):
+		picked_frame = None
+		while picked_frame is None:
+			picked_frame = random.random() * self.frames
+			for skip in self.skip:
+				begin, end = skip
+				if begin <= picked_frame <= end:
+					picked_frame = None
+					break
+		return picked_frame
+	
+	def is_nsfw(self):
+		for nsfw in self.nsfw:
+			begin, end = nsfw
+			if begin <= self.picked_frame <= end:
+				return True
+		return False
+
+	def _translate_basename(self, translations):
+		output_name_tr = self.output_name
+		for k,v in translations.items():
+			replace_token = "${" + k + "}"
+			output_name_tr = output_name_tr.replace(replace_token, v)
+		return output_name_tr
+
+	# TODO: Add translation keywords to the message
+	def get_message(self):
+		return self.message_nsfw if self.is_nsfw() else self.message
+
+	def _render_frame(self):
+		args = [
+			self.render_script,
+			self.path,
+			self.output_name_tr,
+			str(self.picked_frame)
+		]
+		subprocess.run(args)
+
+
+class YandereBot(FediBot.YandereBot):
+	def __init__(self, cfg, keyfile=None, debug_mode=False):
+		settings = {
+				"settings_time": {},
+				"settings_post": {},
+				}
+		self.settings.update(settings)
+		super(YandereBot, self).__init__(cfg, keyfile, debug_mode)
+
+	# Maybe I should remove this from the backend?
+	def print_header_stats(self, picked):
+		profile, frame, nsfw, path = None, None, None, None
+		if picked:
+			profile = picked.profile_name
+			frame = picked.picked_frame
+			nsfw = picked.is_nsfw()
+			path = picked.output_name_tr
+		print("Profile: {} | Frame: {} | NSFW: {} | Path: {}".format(
+			profile, frame, nsfw, path
+			))
+
+
+	def pick(self):
+			dt_picked = self.settings["settings_time"]["datetime"]
+			picked_profile = random.choice(self.settings["settings_post"])
+
+			picked = PickAndRenderFrame(picked_profile, dt_picked)
+			media_list = [picked.output_name_tr]
+			spoiler = picked.is_nsfw()
+			message = picked.get_message()
+			return {
+					"picked":		picked,
+					"media_list":	media_list,
+					"spoiler":		spoiler,
+					"message":		message
+					}
+	
+	def after_pick(self, picked):
+		self.print_header_stats(picked["picked"])
+
+	def post(self, picked):
+		try:
+			super(YandereBot, self).post(picked)
+		except (FileNotFoundError, MastodonAPIError, FediBot.InvalidPost, Exception) as e:
+			print("Exception:", e)
+
+		self.handle_post_exception()
+
+		# The post failed
+		return None
+
 
 
 class FailedToLoadCfg(Exception):
@@ -57,7 +201,7 @@ def main():
 	# Flag if the bot is running in debug mode
 	debug_mode = arguments.dry_run or arguments.debug
 
-	yandere = yandere_bot.YandereBot(
+	yandere = YandereBot(
 		yandere_config,
 		arguments.keyfile,
 		debug_mode
@@ -84,9 +228,9 @@ if __name__ == "__main__":
 	except FailedToLoadCfg:
 		sys.exit(10)
 	# Exceptions raised from the bot
-	except yandere_bot.Debug:
+	except FediBot.Debug:
 		sys.exit(6)
-	except yandere_bot.BadCfgFile:
+	except FediBot.BadCfgFile:
 		sys.exit(4)
-	except yandere_bot.FailedLogin:
+	except FediBot.FailedLogin:
 		sys.exit(2)