#!/usr/bin/env python3

import re
import json
import socketio

WS_URL = "wss://trainmap.belgiantrain.be/socket.io/"


class VERBATIM_KEY: pass
class VERBATIM_VALUE: pass
class DISCARD: pass
def transform_dict(specification: dict, d: dict):
	"""
	specification is a dict where the key is used to determine which action to take with the same
	key in the data. The values to use in specification are:
		- keys that should be discarded: DISCARD
		- keys whose key and value should be retained verbatim: not present
		- keys whose key and/or value should be transformed, with None denoting no change: tuple of
		  (str, function), or any of those two VERBATIM_KEY or VERBATIM_VALUE respectively

	{
	"""
	return {
		(
			specification[k][0]
			if specification.get(k) not in (None, DISCARD) and specification[k][0] is not VERBATIM_KEY
			else k
		): (
			specification[k][1](v)
			if specification.get(k) not in (None, DISCARD) and specification[k][1] is not VERBATIM_VALUE
			else v
		)
		for k, v in d.items()
		if specification.get(k) is not DISCARD
	}


def camel_to_snake(s: str):
	if s.islower():
		return s
	return re.sub(
		r"(?<!^)([A-Z])", # Match an uppercase letter that is not at the start
		r"_\1", # and insert an underscore
		s
	).lower()


def camel_to_snake_keys(data: dict):
	return {
		camel_to_snake(k): v
		for k, v in data.items()
	}


def filtered_data(data):
	return {
		"last_update_success": data["last_updated"]["success"],
		"timestamp": data["timestamp"],
		"trips": {
			trip["trip_short_name"]: transform_dict({
				"trip_id": DISCARD,
				"trip_short_name": DISCARD,
				"position": (VERBATIM_KEY, lambda p: (
					[round(p[0], 6), round(p[1], 6)]
					if isinstance(p, list) else p
				)),
				"current_traveled_distance": (VERBATIM_KEY, lambda x: (
					round(x, 1)
					if isinstance(x, float) else x
				)),
				"fromstationdistance": ("from_station_distance", VERBATIM_VALUE),
				"nextstationdistance": ("next_station_distance", VERBATIM_VALUE),
				"trip_headsign": (VERBATIM_KEY, lambda x: transform_dict({
					"default": DISCARD,
					"en": DISCARD,
					"de": DISCARD,
				}, x)),
			}, camel_to_snake_keys(trip)) if isinstance(trip, dict) else trip
			for trip in data["trips"].values()
		},
	}


def emit(data):
	print(json.dumps(filtered_data(data), ensure_ascii=False))


def main():
	sio = socketio.Client()

	@sio.on("event")
	def handle(data):
		emit(data)

		sio.disconnect()



	sio.connect(WS_URL, transports=["websocket"])

	sio.wait()

if __name__ == "__main__":
	main()