#!/usr/bin/env python3

import argparse
import asyncio
import websockets
import json
import random
import sys

parser = argparse.ArgumentParser(description='Run the taiko-web multiplayer server.')
parser.add_argument('port', type=int, metavar='PORT', nargs='?', default=34802, help='Port to listen on.')
parser.add_argument('-b', '--bind-address', default='localhost', help='Bind server to address.')
parser.add_argument('-o', '--allow-origin', action='append', help='Limit incoming connections to the specified origin. Can be specified multiple times.')
args = parser.parse_args()

server_status = {
	"waiting": {},
	"users": [],
	"invites": {}
}
consonants = "bcdfghjklmnpqrstvwxyz"

def msgobj(msg_type, value=None):
	if value == None:
		return json.dumps({"type": msg_type})
	else:
		return json.dumps({"type": msg_type, "value": value})

def status_event():
	value = []
	for id, userDiff in server_status["waiting"].items():
		value.append({
			"id": id,
			"diff": userDiff["diff"]
		})
	return msgobj("users", value)

def get_invite():
	return "".join([random.choice(consonants) for x in range(5)])

async def notify_status():
	ready_users = [user for user in server_status["users"] if "ws" in user and user["action"] == "ready"]
	if ready_users:
		sent_msg = status_event()
		await asyncio.wait([user["ws"].send(sent_msg) for user in ready_users])

async def connection(ws, path):
	# User connected
	user = {
		"ws": ws,
		"action": "ready",
		"session": False,
		"name": None,
		"don": None
	}
	server_status["users"].append(user)
	try:
		# Notify user about other users
		await ws.send(status_event())
		while True:
			try:
				message = await asyncio.wait_for(ws.recv(), timeout=10)
			except asyncio.TimeoutError:
				# Keep user connected
				pong_waiter = await ws.ping()
				try:
					await asyncio.wait_for(pong_waiter, timeout=10)
				except asyncio.TimeoutError:
					# Disconnect
					break
			except websockets.exceptions.ConnectionClosed:
				# Connection closed
				break
			else:
				# Message received
				try:
					data = json.loads(message)
				except json.decoder.JSONDecodeError:
					data = {}
				action = user["action"]
				msg_type = data["type"] if "type" in data else None
				value = data["value"] if "value" in data else None
				if action == "ready":
					# Not playing or waiting
					if msg_type == "join":
						if value == None:
							continue
						waiting = server_status["waiting"]
						id = value["id"] if "id" in value else None
						diff = value["diff"] if "diff" in value else None
						user["name"] = value["name"] if "name" in value else None
						user["don"] = value["don"] if "don" in value else None
						if not id or not diff:
							continue
						if id not in waiting:
							# Wait for another user
							user["action"] = "waiting"
							user["gameid"] = id
							waiting[id] = {
								"user": user,
								"diff": diff
							}
							await ws.send(msgobj("waiting"))
						else:
							# Join the other user and start game
							user["name"] = value["name"] if "name" in value else None
							user["don"] = value["don"] if "don" in value else None
							user["other_user"] = waiting[id]["user"]
							waiting_diff = waiting[id]["diff"]
							del waiting[id]
							if "ws" in user["other_user"]:
								user["action"] = "loading"
								user["other_user"]["action"] = "loading"
								user["other_user"]["other_user"] = user
								user["other_user"]["player"] = 1
								user["player"] = 2
								await asyncio.wait([
									ws.send(msgobj("gameload", {"diff": waiting_diff, "player": 2})),
									user["other_user"]["ws"].send(msgobj("gameload", {"diff": diff, "player": 1})),
									ws.send(msgobj("name", {
										"name": user["other_user"]["name"],
										"don": user["other_user"]["don"]
									})),
									user["other_user"]["ws"].send(msgobj("name", {
										"name": user["name"],
										"don": user["don"]
									}))
								])
							else:
								# Wait for another user
								del user["other_user"]
								user["action"] = "waiting"
								user["gameid"] = id
								waiting[id] = {
									"user": user,
									"diff": diff
								}
								await ws.send(msgobj("waiting"))
						# Update others on waiting players
						await notify_status()
					elif msg_type == "invite":
						if value and "id" in value and value["id"] == None:
							# Session invite link requested
							invite = get_invite()
							server_status["invites"][invite] = user
							user["action"] = "invite"
							user["session"] = invite
							user["name"] = value["name"] if "name" in value else None
							user["don"] = value["don"] if "don" in value else None
							await ws.send(msgobj("invite", invite))
						elif value and "id" in value and value["id"] in server_status["invites"]:
							# Join a session with the other user
							user["name"] = value["name"] if "name" in value else None
							user["don"] = value["don"] if "don" in value else None
							user["other_user"] = server_status["invites"][value["id"]]
							del server_status["invites"][value["id"]]
							if "ws" in user["other_user"]:
								user["other_user"]["other_user"] = user
								user["action"] = "invite"
								user["session"] = value["id"]
								user["other_user"]["player"] = 1
								user["player"] = 2
								await asyncio.wait([
									ws.send(msgobj("session", {"player": 2})),
									user["other_user"]["ws"].send(msgobj("session", {"player": 1})),
									ws.send(msgobj("invite")),
									ws.send(msgobj("name", {
										"name": user["other_user"]["name"],
										"don": user["other_user"]["don"]
									})),
									user["other_user"]["ws"].send(msgobj("name", {
										"name": user["name"],
										"don": user["don"]
									}))
								])
							else:
								del user["other_user"]
								await ws.send(msgobj("gameend"))
						else:
							# Session code is invalid
							await ws.send(msgobj("gameend"))
				elif action == "waiting" or action == "loading" or action == "loaded":
					# Waiting for another user
					if msg_type == "leave":
						# Stop waiting
						if user["session"]:
							if "other_user" in user and "ws" in user["other_user"]:
								user["action"] = "songsel"
								await asyncio.wait([
									ws.send(msgobj("left")),
									user["other_user"]["ws"].send(msgobj("users", []))
								])
							else:
								user["action"] = "ready"
								user["session"] = False
								await asyncio.wait([
									ws.send(msgobj("gameend")),
									ws.send(status_event())
								])
						else:
							del server_status["waiting"][user["gameid"]]
							del user["gameid"]
							user["action"] = "ready"
							await asyncio.wait([
								ws.send(msgobj("left")),
								notify_status()
							])
					if action == "loading":
						if msg_type == "gamestart":
							user["action"] = "loaded"
							if user["other_user"]["action"] == "loaded":
								user["action"] = "playing"
								user["other_user"]["action"] = "playing"
								sent_msg = msgobj("gamestart")
								await asyncio.wait([
									ws.send(sent_msg),
									user["other_user"]["ws"].send(sent_msg)
								])
				elif action == "playing":
					# Playing with another user
					if "other_user" in user and "ws" in user["other_user"]:
						if msg_type == "note"\
							or msg_type == "drumroll"\
							or msg_type == "branch"\
							or msg_type == "gameresults":
							await user["other_user"]["ws"].send(msgobj(msg_type, value))
						elif msg_type == "songsel" and user["session"]:
							user["action"] = "songsel"
							user["other_user"]["action"] = "songsel"
							sent_msg1 = msgobj("songsel")
							sent_msg2 = msgobj("users", [])
							await asyncio.wait([
								ws.send(sent_msg1),
								ws.send(sent_msg2),
								user["other_user"]["ws"].send(sent_msg1),
								user["other_user"]["ws"].send(sent_msg2)
							])
						elif msg_type == "gameend":
							# User wants to disconnect
							user["action"] = "ready"
							user["other_user"]["action"] = "ready"
							sent_msg1 = msgobj("gameend")
							sent_msg2 = status_event()
							await asyncio.wait([
								ws.send(sent_msg1),
								ws.send(sent_msg2),
								user["other_user"]["ws"].send(sent_msg1),
								user["other_user"]["ws"].send(sent_msg2)
							])
							del user["other_user"]["other_user"]
							del user["other_user"]
					else:
						# Other user disconnected
						user["action"] = "ready"
						user["session"] = False
						await asyncio.wait([
							ws.send(msgobj("gameend")),
							ws.send(status_event())
						])
				elif action == "invite":
					if msg_type == "leave":
						# Cancel session invite
						if user["session"] in server_status["invites"]:
							del server_status["invites"][user["session"]]
						user["action"] = "ready"
						user["session"] = False
						if "other_user" in user and "ws" in user["other_user"]:
							user["other_user"]["action"] = "ready"
							user["other_user"]["session"] = False
							sent_msg = status_event()
							await asyncio.wait([
								ws.send(msgobj("left")),
								ws.send(sent_msg),
								user["other_user"]["ws"].send(msgobj("gameend")),
								user["other_user"]["ws"].send(sent_msg)
							])
						else:
							await asyncio.wait([
								ws.send(msgobj("left")),
								ws.send(status_event())
							])
					elif msg_type == "songsel" and "other_user" in user:
						if "ws" in user["other_user"]:
							user["action"] = "songsel"
							user["other_user"]["action"] = "songsel"
							sent_msg = msgobj(msg_type)
							await asyncio.wait([
								ws.send(sent_msg),
								user["other_user"]["ws"].send(sent_msg)
							])
						else:
							user["action"] = "ready"
							user["session"] = False
							await asyncio.wait([
								ws.send(msgobj("gameend")),
								ws.send(status_event())
							])
				elif action == "songsel":
					# Session song selection
					if "other_user" in user and "ws" in user["other_user"]:
						if msg_type == "songsel" or msg_type == "catjump":
							# Change song select position
							if user["other_user"]["action"] == "songsel" and type(value) is dict:
								value["player"] = user["player"]
								sent_msg = msgobj(msg_type, value)
								await asyncio.wait([
									ws.send(sent_msg),
									user["other_user"]["ws"].send(sent_msg)
								])
						elif msg_type == "crowns" or msg_type == "getcrowns":
							if user["other_user"]["action"] == "songsel":
								sent_msg = msgobj(msg_type, value)
								await asyncio.wait([
									user["other_user"]["ws"].send(sent_msg)
								])
						elif msg_type == "join":
							# Start game
							if value == None:
								continue
							id = value["id"] if "id" in value else None
							diff = value["diff"] if "diff" in value else None
							if not id or not diff:
								continue
							if user["other_user"]["action"] == "waiting":
								user["action"] = "loading"
								user["other_user"]["action"] = "loading"
								await asyncio.wait([
									ws.send(msgobj("gameload", {"diff": user["other_user"]["gamediff"]})),
									user["other_user"]["ws"].send(msgobj("gameload", {"diff": diff}))
								])
							else:
								user["action"] = "waiting"
								user["gamediff"] = diff
								await user["other_user"]["ws"].send(msgobj("users", [{
									"id": id,
									"diff": diff
								}]))
						elif msg_type == "gameend":
							# User wants to disconnect
							user["action"] = "ready"
							user["session"] = False
							user["other_user"]["action"] = "ready"
							user["other_user"]["session"] = False
							sent_msg1 = msgobj("gameend")
							sent_msg2 = status_event()
							await asyncio.wait([
								ws.send(sent_msg1),
								user["other_user"]["ws"].send(sent_msg1)
							])
							await asyncio.wait([
								ws.send(sent_msg2),
								user["other_user"]["ws"].send(sent_msg2)
							])
							del user["other_user"]["other_user"]
							del user["other_user"]
					else:
						# Other user disconnected
						user["action"] = "ready"
						user["session"] = False
						await asyncio.wait([
							ws.send(msgobj("gameend")),
							ws.send(status_event())
						])
	finally:
		# User disconnected
		del user["ws"]
		del server_status["users"][server_status["users"].index(user)]
		if "other_user" in user and "ws" in user["other_user"]:
			user["other_user"]["action"] = "ready"
			user["other_user"]["session"] = False
			await asyncio.wait([
				user["other_user"]["ws"].send(msgobj("gameend")),
				user["other_user"]["ws"].send(status_event())
			])
			del user["other_user"]["other_user"]
		if user["action"] == "waiting":
			del server_status["waiting"][user["gameid"]]
			await notify_status()
		elif user["action"] == "invite" and user["session"] in server_status["invites"]:
			del server_status["invites"][user["session"]]

port = args.port
print('Starting server on port %d' % port)
loop = asyncio.get_event_loop()
tasks = asyncio.gather(
	websockets.serve(connection, args.bind_address, port, origins=args.allow_origin)
)
try:
	loop.run_until_complete(tasks)
	loop.run_forever()
except KeyboardInterrupt:
	print("Stopping server")
	def shutdown_exception_handler(loop, context):
		if "exception" not in context or not isinstance(context["exception"], asyncio.CancelledError):
			loop.default_exception_handler(context)
	loop.set_exception_handler(shutdown_exception_handler)
	tasks = asyncio.gather(*asyncio.all_tasks(loop=loop), loop=loop, return_exceptions=True)
	tasks.add_done_callback(lambda t: loop.stop())
	tasks.cancel()
	while not tasks.done() and not loop.is_closed():
		loop.run_forever()
finally:
	if hasattr(loop, "shutdown_asyncgens"):
		loop.run_until_complete(loop.shutdown_asyncgens())
	loop.close()