#! /usr/bin/env python3

import socket, struct, dataclasses

@dataclasses.dataclass
class Game:
	version: int # ui32
	name: str # 64
	size: int # i32
	alliances: int # ui8
	tech_level: int # ui8
	power_level: int # ui8
	bases_level: int # ui8
	host_address: str # 40
	max_players: int # i32
	current_players: int # i32
	game_type: int # ui32
	current_spectators: int # ui16
	max_spectators: int # ui16
	blind_mode: int # ui32
	unused1: int # ui32
	host2: str # 40
	host3: str # 40
	unused2: str # 157
	host_port: int # ui16
	map_name: str # 40
	host_name: str # 40
	version_string: str # 64
	mod_list: str # 255
	version_major: int # ui32
	version_minor: int # ui32
	private_game: int # ui32
	pure_map: int # ui32
	mods: int # ui32
	game_id: int # ui32
	limits: int # ui32
	os: int # ui32
	version_major_net: int # ui16
	version_minor_net: int # ui16

	binary_format = struct.Struct('!I64si4B40s2iI2H2I40s40s157sH40s40s64s255s8I2H')

	def to_bytes(self):
		values = []
		for field in dataclasses.fields(__class__):
			value = getattr(self, field.name)
			if field.type is str: value = value.encode()
			values.append(value)
		return __class__.binary_format.pack(*values)

	@staticmethod
	def from_bytes(bytes): return Game(*__class__.binary_format.unpack(bytes))

	def __post_init__(self):
		for field in dataclasses.fields(__class__):
			if field.type is str: 
				setattr(self, field.name, getattr(self, field.name).split(b'\0', 1)[0].decode())

def receive_list(host, port, source_address='', query_command=b'list\0'):
	with socket.create_connection(
		(host, port),
		5,
		(source_address, 0)
	) as s:
		def send(binary_format, *data): s.sendall(struct.pack(binary_format, *data))
		
		def receive(binary_format):
			binary_format = struct.Struct(binary_format)
			received = bytearray()
			size_left = binary_format.size
			while size_left:
				chunk = s.recv(size_left)
				if not chunk: raise EOFError('socket connection broken')
				received += chunk
				size_left -= len(chunk)
			return binary_format.unpack(received)

		send('5s', query_command)
		
		count, = receive('!I')
		games = []
		for _ in range(count): games.append(Game(*receive(Game.binary_format.format)))

		status, = receive('!I')

		motd_len, = receive('!I')
		motd, = receive(f'!{motd_len}s')
		motd = motd.decode()

		more_games, = receive('!I')
		if more_games & 1: # See explanation in warzone2100/lib/netplay/netplay.cpp function NETenumerateGames 
			count, = receive('!I')
			games = []
			for _ in range(count): games.append(Game(*receive(Game.binary_format.format)))

		return status, motd, games

def main():
		import argparse
		parser = argparse.ArgumentParser()
		parser.add_argument('--source-address', nargs='?', default='')
		parser.add_argument('--output-format', choices=['text', 'json', 'repr'], default='text')
		parser.add_argument('host', nargs='?', default='warzone2100.retropaganda.info')
		parser.add_argument('port', nargs='?', type=int, default=9990)
		args = parser.parse_args()

		status, motd, games = receive_list(args.host, args.port, args.source_address)

		if args.output_format == 'text':
			print(f'Status: {status}')
			print()
			print('MotD:')
			print('\t' + '\n\t'.join(motd.split('\n')))
			print()
			print(f'Game count: {len(games)}')
			print()
			print('Confederate address              Port   GId  Host address      Port  OS       Spect  Plays  Map              Host name        Title')
			for game in games:
				os = {  0x6c696e: 'UNIX', 0x6d6163: 'MacOSX', 0x77696e: 'Windows', 1001: '-' }.get(game.os, str(game.os))
				print(f'{game.host2:30.30}  {game.host3:>5}  {game.game_id:>4}  {game.host_address:15.15}  {game.host_port:>5}  {os:7.7}  ', end='')
				print(f'{f'{game.current_spectators:>2}' if game.current_spectators else '  '}/{game.max_spectators:>2}  ', end='')
				print(f'{f'{game.current_players:>2}' if game.current_players != 0 else '  '}/{game.max_players:>2}  ', end='')
				print(f'{game.map_name:15.15}  {game.host_name:15.15}  {game.name}')
			print()
		elif args.output_format == 'json':
			import json
			print(json.dumps(
				{
					'status': status,
					'motd': motd,
					'games': [dataclasses.asdict(game) for game in games]
				},
				check_circular=False, ensure_ascii=False, indent='\t'
			))
		else: # args.output_format == 'repr'
			for game in games:
				print(game)
				print()

if __name__ == '__main__': main()