Post_DvZ_Bot/servers.py

import pickle
import os
import math
import asyncio
from mcstatus import MinecraftServer
import datetime
#we use the ips of mc servers like ids of discord servers
#Important note: DCServer objects must be created before MCServer objects

client = None

#these dictionaries store objects
# with discord server ids/mc server ips as keys
MCServers = {}
DCServers = {}

#this dictionaries store sets of mc server ips with
# discord server id as keys and the other way around.
linksMC_DC = {}
linksDC_MC = {}

def link(mc_ip, dc_id):
	try:
		linksMC_DC[mc_ip] = linksMC_DC[mc_ip].union([dc_id])
	except KeyError:
		linksMC_DC[mc_ip] = set([dc_id])
	try:
		linksDC_MC[dc_id] = linksDC_MC[dc_id].union([mc_ip])
	except KeyError:
		linksDC_MC[dc_id] = set([mc_ip])

def get_time_for_next_update(count, threshold):
	if count < threshold / 2:
		return 15
	if count >= threshold:
		return 1
	return 5

class MCServer:
	def __init__(self, ip, name, game_threshold, DCs):
		self.ip = ip
		self.name = name
		self.game_threshold = game_threshold
		MCServers[ip] = self
		self.time_since_last_update = math.inf
		self.count = 0
		for id in DCs:
			link(ip, id)
	#runs once per minute
	def update_tick(self):
		self.time_since_last_update += 1
		if self.time_since_last_update >= get_time_for_next_update(self.count, self.game_threshold):
			self.time_since_last_update = 0
			return self
		else:
			return None
	def set_count(self, count):
		changed = self.count == count
		self.count = count
		return changed
		
			

class DCServer:
	def __init__(self, id, channelID):
		self.id = id
		self.channelID = channelID
		DCServers[id] = self
	def get_player_count_string(self):
		countstr = "Players Online: "
		for mc in linksDC_MC[self.id]:
			mc = MCServers[mc]
			if mc.count == -1:
				countstr += "Can't reach " + mc.name
			else:
				countstr += mc.name + " " + str(mc.count)
				if mc.count >= mc.game_threshold:
					countstr += "✳"
			countstr += " +++ "
		return countstr
		

#Read servers and links from file if file exists
try:
	dirname, filename = os.path.split(os.path.abspath(__file__))
	with open(os.path.join(dirname, "servers.ref"), "rb") as f:
		MCServers, DCServers, linksMC_DC, linksDC_MC = pickle.load(f)
		for server in MCServers:
			#update all servers on start
			MCServers[server].time_since_last_update = math.inf
except FileNotFoundError:
	pass #use empty dictionaries as defined earlier


#creates DCServer object if it doesn't exist yet
def ensure_discord(id, channelID):
	try:
		DCServers[id] #already exists
	except KeyError:
		#create new DCServer object,
		# gets put into dict in constructor
		DCServer(id, channelID)

def parse_arguments(command_content):
	words = command_content.split()
	return words[1], words[2], int(words[3])

	
#returns response to message
def add_mc_server(command):
	try:
		ip, name, threshold = parse_arguments(command.content)
	except IndexError:
		return "Could not instantiate Server object: Too few arguments"
	except ValueError:
		return "Could not instantiate Server object: The 3rd argument must be an integer"
		
	try:
		ensure_discord(command.guild.id, command.channel.id)
	except AttributeError:
		return "This command does not work in DM channels"
	try:
		MCServers[ip]
		#mc server already exists, link with discord
		try:
			if command.guild.id in linksMC_DC[ip]:
				return "already linked" #maybe make this unlink
			else:
				link(ip, command.guild.id)
		except KeyError:
			link(ip, command.guild.id)
	except KeyError:
		#create new mc server object
		MCServer(ip, name, threshold, [command.guild.id])
	#save servers to file
	dirname, filename = os.path.split(os.path.abspath(__file__))
	with open(os.path.join(dirname, "servers.ref"), "wb") as f:
		pickle.dump((MCServers, DCServers, linksMC_DC, linksDC_MC), f)
	return "Successfully linked MC Server"


class Updater:
	mcServersToUpdate = set()
	dcServersToUpdate = set()
	force = False
	async def force_update(self, discordID):
		self.force = True
		mcs = linksDC_MC[discordID]
		for mc in mcs:
			mc = MCServers[mc]
			mc.time_since_last_update = math.inf
		self.update_tick()
		await self.wait_for_force()
		return DCServers[discordID].get_player_count_string()
	async def wait_for_force(self):
		while self.force:
			await asyncio.sleep(1)
	async def start(self):
		while True:
			now = datetime.datetime.now()
			after_minute = now.second + now.microsecond / 1_000_000
			if after_minute:
				await asyncio.sleep(60 - after_minute)
			self.update_tick()
	
	#called when all player counts are retrieved
	def done_updating(self):
		for update in self.updateResults:
			should_update_header, mcserver, _ = update.result()
			if should_update_header:
				self.dcServersToUpdate = self.dcServersToUpdate.union(linksMC_DC[mcserver.ip])		
		for dc in self.dcServersToUpdate:
			dc = DCServers[dc]
			channel = client.get_channel(dc.channelID)
			if channel != None:
				loop = asyncio.get_event_loop()
				loop.create_task(channel.edit(topic = dc.get_player_count_string()))
		print("Updated player counter:", datetime.datetime.now())
		self.force = False
		#reSET these
		self.mcServersToUpdate = set()
		self.dcServersToUpdate = set()
	
	#called on the start of every minute
	def update_tick(self):
		for name in MCServers:
			s = MCServers[name].update_tick()
			if s != None:
				self.mcServersToUpdate = self.mcServersToUpdate.union([s])
		
		self.updatesToDo = len(self.mcServersToUpdate)
		self.updateResults = set()
		loop = asyncio.get_event_loop()
		
		def on_done(future):
			a, b, updater = future.result()
			updater.updatesToDo -= 1
			if updater.updatesToDo < 1:
				updater.done_updating()
		
		for server in self.mcServersToUpdate:
			f = loop.create_future()
			asyncio.ensure_future(self.update_server(f, server))
			f.add_done_callback(on_done)
			self.updateResults = self.updateResults.union([f])
		
	async def update_server(self, future, server):
		try:
			m_server = MinecraftServer.lookup(server.ip)
			count = m_server.status().players.online
			future.set_result((server.set_count(count), server, self))
		except:
			future.set_result((server.set_count(-1), server, self))