MIND/backend/server.py

#-*- coding: utf-8 -*-

from __future__ import annotations

from os import execv, urandom
from sys import argv
from threading import Timer, current_thread
from typing import TYPE_CHECKING, List, NoReturn, Union

from flask import Flask, render_template, request
from waitress import create_server
from waitress.task import ThreadedTaskDispatcher as TTD
from werkzeug.middleware.dispatcher import DispatcherMiddleware

from backend.db import DB_Singleton, DBConnection, close_db, revert_db_import
from backend.helpers import RestartVars, Singleton, folder_path
from backend.logging import LOGGER
from backend.settings import restore_hosting_settings

if TYPE_CHECKING:
	from waitress.server import TcpWSGIServer

THREADS = 10

class ThreadedTaskDispatcher(TTD):
	def handler_thread(self, thread_no: int) -> None:
		super().handler_thread(thread_no)
		i = f'{DBConnection}{current_thread()}'
		if i in DB_Singleton._instances and not DB_Singleton._instances[i].closed:
			DB_Singleton._instances[i].close()
		return

	def shutdown(self, cancel_pending: bool = True, timeout: int = 5) -> bool:
		print()
		LOGGER.info('Shutting down MIND')
		result = super().shutdown(cancel_pending, timeout)
		DBConnection(timeout=20.0).close()
		return result


class Server(metaclass=Singleton):
	api_prefix = "/api"
	admin_api_extension = "/admin"
	admin_prefix = "/api/admin"

	def __init__(self) -> None:
		self.do_restart = False
		"Restart instead of shutdown"

		self.restart_args: List[str] = []
		"Flag to run with when restarting"

		self.handle_flags: bool = False
		"Run any flag specific actions before restarting"

		self.url_prefix = ""

		self.revert_db_timer = Timer(60.0, self.__revert_db)
		self.revert_db_timer.name = "DatabaseImportHandler"
		self.revert_hosting_timer = Timer(60.0, self.__revert_hosting)
		self.revert_hosting_timer.name = "HostingHandler"

		return

	def create_app(self) -> None:
		"""Create a Flask app instance"""
		from frontend.api import admin_api, api
		from frontend.ui import ui

		app = Flask(
			__name__,
			template_folder=folder_path('frontend','templates'),
			static_folder=folder_path('frontend','static'),
			static_url_path='/static'
		)
		app.config['SECRET_KEY'] = urandom(32)
		app.config['JSONIFY_PRETTYPRINT_REGULAR'] = True
		app.config['JSON_SORT_KEYS'] = False

		# Add error handlers
		@app.errorhandler(400)
		def bad_request(e):
			return {'error': 'Bad request', 'result': {}}, 400

		@app.errorhandler(405)
		def method_not_allowed(e):
			return {'error': 'Method not allowed', 'result': {}}, 405

		@app.errorhandler(500)
		def internal_error(e):
			return {'error': 'Internal error', 'result': {}}, 500

		@app.errorhandler(404)
		def not_found(e):
			if request.path.startswith(self.api_prefix):
				return {'error': 'Not Found', 'result': {}}, 404
			return render_template('page_not_found.html', url_prefix=self.url_prefix)

		app.register_blueprint(ui)
		app.register_blueprint(api, url_prefix=self.api_prefix)
		app.register_blueprint(admin_api, url_prefix=self.admin_prefix)

		# Setup closing database
		app.teardown_appcontext(close_db)

		self.app = app
		return

	def set_url_prefix(self, url_prefix: str) -> None:
		"""Change the URL prefix of the server.

		Args:
			url_prefix (str): The desired URL prefix to set it to.
		"""
		self.app.config["APPLICATION_ROOT"] = url_prefix
		self.app.wsgi_app = DispatcherMiddleware(
			Flask(__name__),
			{url_prefix: self.app.wsgi_app}
		)
		self.url_prefix = url_prefix
		return

	def __create_waitress_server(
		self,
		host: str,
		port: int
	) -> TcpWSGIServer:
		"""From the `Flask` instance created in `self.create_app()`, create
		a waitress server instance.

		Args:
			host (str): The host to bind to.
			port (int): The port to listen on.

		Returns:
			TcpWSGIServer: The waitress server.
		"""
		dispatcher = ThreadedTaskDispatcher()
		dispatcher.set_thread_count(THREADS)
		server = create_server(
			self.app,
			_dispatcher=dispatcher,
			host=host,
			port=port,
			threads=THREADS
		)
		return server

	def run(self, host: str, port: int) -> None:
		"""Start the webserver.

		Args:
			host (str): The host to bind to.
			port (int): The port to listen on.
		"""
		self.server = self.__create_waitress_server(host, port)
		LOGGER.info(f'MIND running on http://{host}:{port}{self.url_prefix}')
		self.server.run()

		return

	def __shutdown_thread_function(self) -> None:
		"""Shutdown waitress server. Intended to be run in a thread.
		"""
		self.server.close()
		self.server.task_dispatcher.shutdown()
		self.server._map.clear()
		return

	def shutdown(self) -> None:
		"""Stop the waitress server. Starts a thread that
		shuts down the server.
		"""
		t = Timer(1.0, self.__shutdown_thread_function)
		t.name = "InternalStateHandler"
		t.start()
		return

	def restart(
		self,
		restart_args: List[str] = [],
		handle_flags: bool = False
	) -> None:
		"""Same as `self.shutdown()`, but restart instead of shutting down.

		Args:
			restart_args (List[str], optional): Any arguments to run the new instance with.
				Defaults to [].

			handle_flags (bool, optional): Run flag specific actions just before restarting.
				Defaults to False.
		"""		
		self.do_restart = True
		self.restart_args = restart_args
		self.handle_flags = handle_flags
		self.shutdown()
		return

	def handle_restart(self, flag: Union[str, None]) -> NoReturn:
		"""Restart the interpreter.

		Args:
			flag (Union[str, None]): Supplied flag, for flag handling.

		Returns:
			NoReturn: No return because it replaces the interpreter.
		"""
		if self.handle_flags:
			handle_flags_pre_restart(flag)

		LOGGER.info('Restarting MIND')
		from MIND import __file__ as mind_file
		execv(folder_path(mind_file), [argv[0], *self.restart_args])

	def __revert_db(self) -> None:
		"""Revert database import and restart.
		"""
		LOGGER.warning(f'Timer for database import expired; reverting back to original file')
		self.restart(handle_flags=True)
		return

	def __revert_hosting(self) -> None:
		"""Revert the hosting changes.
		"""
		LOGGER.warning(f'Timer for hosting changes expired; reverting back to original settings')
		self.restart(handle_flags=True)
		return


SERVER = Server()


def handle_flags(flag: Union[None, str]) -> None:
	"""Run flag specific actions on startup.

	Args:
		flag (Union[None, str]): The flag or `None` if there is no flag set.
	"""
	if flag == RestartVars.DB_IMPORT:
		LOGGER.info('Starting timer for database import')
		SERVER.revert_db_timer.start()

	elif flag == RestartVars.HOST_CHANGE:
		LOGGER.info('Starting timer for hosting changes')
		SERVER.revert_hosting_timer.start()

	return


def handle_flags_pre_restart(flag: Union[None, str]) -> None:
	"""Run flag specific actions just before restarting.

	Args:
		flag (Union[None, str]): The flag or `None` if there is no flag set.
	"""
	if flag == RestartVars.DB_IMPORT:
		revert_db_import(swap=True)

	elif flag == RestartVars.HOST_CHANGE:
		with SERVER.app.app_context():
			restore_hosting_settings()
			close_db()

	return