#!/usr/bin/env python3

import string
import shutil
import os
import pyuavcan_v0 as uavcan

def get_canros_dir():
	return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

#try:
#	import canros
#except ModuleNotFoundError:
	# If canros hasn't been built before, add the src folder to sys.path and re-import
# TODO ignoring this because this should only ever be used during ROS2 build,
# decide if we ever need to actually import the package proper
# Maybe if we call this script after setuptools I can just import it normally?
import sys
sys.path.append(os.path.join(get_canros_dir(), "canros"))
import uavcan_msgs
import constants
del sys

def print_type(typ, f=None):
	if typ.category == typ.CATEGORY_PRIMITIVE:
		if typ.kind == typ.KIND_BOOLEAN:
			print("bool", end='', file=f)
		elif typ.kind == typ.KIND_FLOAT:
			if typ.bitlen <= 32:
				print("float32", end='', file=f)
			else:
				print("float64", end='', file=f)
		else:
			if typ.kind == typ.KIND_UNSIGNED_INT:
				print("u", end='', file=f)
			if typ.bitlen <= 8:
				print("int8", end='', file=f)
			elif typ.bitlen <= 16:
				print("int16", end='', file=f)
			elif typ.bitlen <= 32:
				print("int32", end='', file=f)
			else:
				print("int64", end='', file=f)

	elif typ.category == typ.CATEGORY_COMPOUND:
		print(uavcan_msgs.ros_name(typ.full_name), end='', file=f)

	elif typ.category == typ.CATEGORY_ARRAY:
		print_type(typ.value_type, f=f)
		print("[]", sep='', end='', file=f)

def print_constant(const, f=None):
	print_type(const.type, f=f)
	print(" ", const.name, "=", const.value, sep='', file=f)

def print_field(field, f=None):
	print_type(field.type, f=f)
	print(" ", field.name, sep='', file=f)

def print_id(f=None):
	id_tag_type = uavcan.dsdl.parser.PrimitiveType(
		uavcan.dsdl.parser.PrimitiveType.KIND_UNSIGNED_INT,
		7,
		uavcan.dsdl.parser.PrimitiveType.CAST_MODE_SATURATED
	)

	print("\n# UAVCAN ID field", file=f)
	print_field(uavcan.dsdl.parser.Field(id_tag_type, constants.uavcan_id_field_name), f=f)

def print_union(fields, f=None):
	union_tag_type = uavcan.dsdl.parser.PrimitiveType(
		uavcan.dsdl.parser.PrimitiveType.KIND_UNSIGNED_INT,
		len(fields),
		uavcan.dsdl.parser.PrimitiveType.CAST_MODE_SATURATED
	)

	print("\n# Union field selector tag", file=f)

	for i in range(0, len(fields)):
		print_constant(uavcan.dsdl.parser.Constant(
			union_tag_type,
			constants.union_const_prefix + fields[i].name.upper(),
			str(i), i
		), f=f)

	print_field(uavcan.dsdl.parser.Field(union_tag_type, constants.union_tag_field_name), f=f)

def print_msg(constants, fields, union, f=None):
	print("# Internal field for canros' usage", file=f)
	print("bool msg_from_canros false", file=f)
	for const in constants:
		if const.type.category == const.type.CATEGORY_VOID:
			continue
		print_constant(const, f=f)

	for field in fields:
		if field.type.category == field.type.CATEGORY_VOID:
			continue
		print_field(field, f=f)

	if union:
		print_union(fields, f=f)

def print_node_info_custom(f=None):
	print_header(["Generated by canros", "Custom NodeInfo message."], f=f)
	print(
		string.Template("""$NodeStatus status
$SoftwareVersion software_version
$HardwareVersion hardware_version
uint8[] name
""").substitute(
			NodeStatus=uavcan_msgs.ros_name("uavcan.protocol.NodeStatus"),
			SoftwareVersion=uavcan_msgs.ros_name("uavcan.protocol.SoftwareVersion"),
			HardwareVersion=uavcan_msgs.ros_name("uavcan.protocol.HardwareVersion")
		),
		sep='', end='', file=f
	)

def print_get_node_info_custom(f=None):
	print_header([
		"Generated by canros",
		"Custom NodeInfo service.",
		"Can be used to obtain configuration of the canros server."
	], f=f)
	print(
		"""---
NodeInfo node_info
""",
		sep='', end='', file=f
	)

def print_header(rows, f=None):
	maxlen = max([len(row) for row in rows])
	horz_bar = "#"*(maxlen+6)
	print(horz_bar, file=f)
	for row in rows:
		print("# ", row, " "*(maxlen-len(row)), "#", file=f)
	print(horz_bar, "\n", sep='', file=f)

def make_rel_dir(directory):
	abs_dir = os.path.join(get_canros_dir(), directory)
	if os.path.exists(abs_dir):
		shutil.rmtree(abs_dir)
	os.makedirs(abs_dir)
	return abs_dir

def main():
	ros_msg_dir = make_rel_dir("msg")
	ros_srv_dir = make_rel_dir("srv")

	for _, typ in uavcan.TYPENAMES.items():
		if typ.kind == typ.KIND_MESSAGE:
			file_path = os.path.join(ros_msg_dir, uavcan_msgs.ros_name(typ.full_name) + ".msg")
		else:
			file_path = os.path.join(ros_srv_dir, uavcan_msgs.ros_name(typ.full_name) + ".srv")

		with open(file_path, 'w+') as f:
			print_header(["Generated by canros", typ.full_name], f=f)

			if typ.kind == typ.KIND_MESSAGE:
				print_msg(typ.constants, typ.fields, typ.union, f=f)
				if uavcan_msgs.ros_has_id_field(typ.full_name):
					print_id(f=f)
			elif typ.kind == typ.KIND_SERVICE:
				print_msg(typ.request_constants, typ.request_fields, typ.request_union, f=f)
				print_id(f=f)
				print("---", file=f)
				print_msg(typ.response_constants, typ.response_fields, typ.response_union, f=f)

	# Custom messages
	with open(os.path.join(ros_msg_dir, "NodeInfo.msg"), 'w+') as f:
		print_node_info_custom(f=f)
	with open(os.path.join(ros_srv_dir, "GetNodeInfo.srv"), 'w+') as f:
		print_get_node_info_custom(f=f)

if __name__ == "__main__":
	main()
