compile

#!/bin/python3

# no empty lines with whitespace
# more comments
# simplify

#comment and format generated code better

import yaml
import sys
from enum import Enum

def get_type_identifier(type):
	switch = {
		'u1': 'uint8_t',
		'u2': 'uint16_t',
		'u4': 'uint32_t',
		's1': 'int8_t',
		's2': 'int16_t',
		's4': 'int32_t',
	}
	return switch.get(type, type)

def split_expression(expr):
	def is_special_char(c):
		return c in [' ', '+', '-', '*', '/', '&', '|', '(', ')', '!', '=', '>', '<']

	pos_split = []
	strings = []
	for i, current in enumerate(expr):
		if i != 0:
			last = expr[i - 1]
			if (is_special_char(current) and (not is_special_char(last))) or ((not is_special_char(current)) and is_special_char(last)):
				pos_split.append(i)
	pos_split.append(len(expr))
	for i, split in enumerate(pos_split):
		start = 0 if i == 0 else pos_split[i - 1]
		str = expr[start : split].strip()
		if len(str) > 0:
			strings.append(str)
	return strings

def make_enum_name(type, name):
	type = type.strip()
	ret = f'{type}_{name}'

	return ret.upper()

def prefix_expression(expr):
	pos_split = []
	parts = split_expression(expr)
	for i, part in enumerate(parts):
		switch = {
			'and': '&&',
			'or': '||',
			'not': '!',
			'true': '1',
			'false': '0',
		}

		res = switch.get(part, None)
		if res != None:
			parts[i] = res
		else:
			char = part[0]
			if char.isalpha():
				parts[i] = f'data->{part}'
			if '::' in part:
				split = part.split('::')
				parts[i] = make_enum_name(split[0], split[1])
	expr = ' '.join(parts)
	expr = expr.replace('_.', '##current##->')
	expr = expr.replace('_root.', '_root->')
	expr = expr.replace('( ', '(')
	expr = expr.replace(' )', ')')
	return expr

def make_expression_for_setting_size(expr, size):
	# Bail on impossible calculations with more than one variable
	if expr.count('data->') > 1:
		return None

	# Bail when multiplication/division is used
	if '*' in expr or '%' in expr or '/' in expr:
		return None

	# Should be 'ret' '->' 'NAME'
	len_expr_left = 3

	split = split_expression(expr)
	if len(split) == len_expr_left:
		return f'{expr} = {size};'

	left_expr = ''.join(split[:len_expr_left])
	right_expr = ' '.join(split[len_expr_left:])
	return f'{left_expr} = {size} - ({right_expr});'


def find_type_with_ok_dependencies(existing, remaining):
	for i, current_to_test in enumerate(remaining):
		all_satisfied = True
		for current_dependency in current_to_test.seq.dependencies:
			found = False
			for ex in existing:
				if ex.name == current_dependency:
					found = True
			if not found:
				all_satisfied = False
				break
		if all_satisfied:
			return i
	return -1

def sort_types(info):
	ret = []
	while len(info) > 0:
		index = find_type_with_ok_dependencies(ret, info)
		if index == -1:
			print('Remaining types:')
			print([(x.name) for x in info])
			raise Exception('Can not find type!')
		ret.append(info[index])
		info.pop(index)

	return ret

def write_struct(sb, name, seq):
	sb.append(f'typedef struct {name}_\n')
	sb.append('{\n')
	for line in seq.lines_struct:
		sb.append(line + '\n')
	sb.append(f'}} {name};\n\n')

class method_type(Enum):
	READ = 1
	WRITE = 2
	FREE = 3

def write_method(sb, name, code, type, type_root):
	if type == method_type.READ:
		sb.append(f'int stream_read_{name}(stream *s_root, stream *s, {name} *data, {type_root} *_root)\n')
	elif type == method_type.WRITE:
		sb.append(f'int stream_write_{name}(stream *s_root, stream *s, {name} *data, stream_write_stage stage, BOOL do_sequence, {type_root} *_root)\n')
	elif type == method_type.FREE:
		sb.append(f'void stream_free_{name}({name} *data)\n')
	sb.append('{\n')
	for line in code.lines_var:
		sb.append(line + '\n')
	if code.has_i:
		sb.append('\tint i;\n')
	if len(code.lines_var) > 0 or code.has_i:
		sb.append('\n')
	for line in code.lines_code:
		sb.append(line + '\n')
	if type != method_type.FREE:
		sb.append('\treturn 0;\n')
	sb.append('}\n\n')

class type_info:
	def __init__(self, name):
		self.name = name
		self.seq = seq_info()

	def load(self, node):
		self.seq.write.add_code('/* No separate sequence run during write */')
		self.seq.write.add_code('if (stage == STREAM_WRITE_STAGE_WRITE && do_sequence) return 0;')
		self.seq.write.add_line()
		for k, v in node.items():
			if k == 'seq':
				self.seq.load_sequences(v)
			if k == 'instances':
				self.seq.load_instances(v)

class enum_info:
	def __init__(self):
		self.name = ''
		self.lines_enum = []

	def add(self, str):
		self.lines_enum.append('\t' + str)

	def load(self, name, node):
		self.name = name
		for k, v in node.items():
			name2 = make_enum_name(name, v)
			self.add(f'{name2} = {k},')


class type_switch_info:
	def __init__(self):
		self.enumname = ''
		self.type = ''
		self.unionname = ''

class code_info:
	def __init__(self):
		self.lines_code = []
		self.lines_var = []
		self.has_i = False
		self.indent_code = 1

	def add_code(self, str):
		tabs = '\t' * self.indent_code
		self.lines_code.append(tabs + str)

	def add_var(self, str):
		self.lines_var.append('\t' + str)

	def indent_code_plus(self):
		self.add_code('{')
		self.indent_code += 1

	def indent_code_minus(self):
		self.indent_code -= 1
		self.add_code('}')

	def add_i(self):
		self.has_i = True

	def add_line(self):
		self.lines_code.append('')

class seq_info:
	def __init__(self):
		self.read = code_info()
		self.write = code_info()
		self.free = code_info()
		self.lines_struct = []
		self.dependencies = []
		self.indent_struct = 1

	def add_code_all(self, str):
		self.read.add_code(str)
		self.write.add_code(str)
		self.free.add_code(str)

	def add_code_readwrite(self, str):
		self.read.add_code(str)
		self.write.add_code(str)

	def add_var_all(self, str):
		self.read.add_var(str)
		self.write.add_var(str)
		self.free.add_var(str)

	def add_var_readwrite(self, str):
		self.read.add_var(str)
		self.write.add_var(str)

	def add_struct(self, str):
		tabs = '\t' * self.indent_struct
		line = tabs + str
		if line not in self.lines_struct:
			self.lines_struct.append(line)

	def add_dependency(self, str):
		if str.startswith('uint') or str.startswith('int'):
			return
		if str not in self.dependencies:
			self.dependencies.append(str)

	def indent_code_plus_all(self):
		self.read.indent_code_plus();
		self.write.indent_code_plus();
		self.free.indent_code_plus();

	def indent_code_plus_readwrite(self):
		self.read.indent_code_plus();
		self.write.indent_code_plus();

	def indent_code_minus_all(self):
		self.read.indent_code_minus();
		self.write.indent_code_minus();
		self.free.indent_code_minus();

	def indent_code_minus_readwrite(self):
		self.read.indent_code_minus();
		self.write.indent_code_minus();

	def add_i(self):
		self.read.add_i()
		self.write.add_i()
		self.free.add_i()

	def indent_struct_plus(self):
		self.add_struct('{')
		self.indent_struct += 1

	def indent_struct_minus(self, str):
		self.indent_struct -= 1
		self.add_struct(str)

	def load_instance(self, instance_name, node, first):
		node_io = None
		node_pos = None
		node_type = None
		node_value = None
		node_condition = None
		node_enum = None

		# Get values from node
		for key, value in node.items():
			if key == 'io':
				node_io = value
			if key == 'pos':
				node_pos = prefix_expression(value)
			if key == 'type':
				node_type = value
			if key == 'value':
				node_value = prefix_expression(value)
			if key == 'if':
				node_condition = prefix_expression(value)
			if key == 'enum':
				node_enum = value

		# Generate code for loading instance

		if node_value is not None:
			if node_condition is not None:
				self.read.add_code(f'if ({node_condition})', )
				self.read.indent_code_plus()

			struct_type = node_enum if node_enum is not None else 'int32_t'

			self.add_struct(f'{struct_type} {instance_name};')
			self.read.add_code(f'data->{instance_name} = {node_value};')

			if node_condition is not None:
				self.read.indent_code_minus()

			return

		if node_io != '_root._io':
			raise Exception('Unsupported')

		if not first:
			self.write.add_line()
		self.write.add_code('/* Store position for instance */')
		self.write.add_code('if (stage == STREAM_WRITE_STAGE_DRYRUN)')
		self.write.indent_code_plus()
		self.add_struct(f'int _dryrun_pos_instance_{instance_name};')
		self.write.add_code(f'data->_dryrun_pos_instance_{instance_name} = stream_write_get_position_absolute(s_root);')
		expr_size = make_expression_for_setting_size(node_pos, 'stream_write_get_position_absolute(s_root)')
		if expr_size is not None:
			self.write.add_code(expr_size)
		self.write.indent_code_minus()

		if node_condition is not None:
			self.add_code_all(f'if ({node_condition})')
			self.indent_code_plus_all()

		self.add_struct(f'struct {node_type}_ *{instance_name};')
		self.read.add_var(f'stream substream_instance_{instance_name};')
		self.write.add_var(f'stream substream_instance_{instance_name};')
		self.read.add_code(f'CHECK(stream_read_make_substream_instance(s_root, &substream_instance_{instance_name}, ({node_pos}), s_root->max - ({node_pos})));')

		if node_condition == None:
			self.write.add_line()
		self.write.add_code('/* Make substream for instance */')
		self.write.add_code(f'if (stage == STREAM_WRITE_STAGE_WRITE)');
		self.write.indent_code_plus()
		self.write.add_code(f'CHECK(stream_write_make_substream_instance(s_root, &substream_instance_{instance_name}, data->_dryrun_pos_instance_{instance_name}));')
		self.write.indent_code_minus()
		self.write.add_code('else')
		self.write.indent_code_plus()
		self.write.add_code(f'CHECK(stream_write_make_substream(s_root, &substream_instance_{instance_name}));')
		self.write.indent_code_minus()

		self.write.add_line()
		self.write.add_code('/* Write instance */')
		if not node_type.startswith('uint') and not node_type.startswith('int'):
			self.read.add_code(f'data->{instance_name} = malloc(sizeof({node_type}));')
		self.read.add_code(f'CHECK(stream_read_{node_type}(s_root, &substream_instance_{instance_name}, data->{instance_name}, _root));')
		self.write.add_code(f'CHECK(stream_write_{node_type}(s_root, &substream_instance_{instance_name}, data->{instance_name}, stage, TRUE, _root));')
		self.write.add_code(f'CHECK(stream_write_{node_type}(s_root, &substream_instance_{instance_name}, data->{instance_name}, stage, FALSE, _root));')
		self.free.add_code(f'stream_free_{node_type}(data->{instance_name});')
		self.free.add_code(f'free(data->{instance_name});')

		if node_condition is not None:
			self.indent_code_minus_all()

	def load_instances(self, node):
		if len(node) > 0:
			self.write.add_line()
			self.write.add_code('/* Start writing instance data */')
			self.write.add_code('if (!do_sequence)')
			self.write.indent_code_plus()
		first = True
		for name, value in node.items():
			self.load_instance(name, value, first)
			first = False
		if len(node) > 0:
			self.write.indent_code_minus()

	def load_sequence(self, child):
		node_id = None
		node_contents = None
		node_type = None
		node_switch_on = None
		node_type_switch = []
		node_repeat = None
		node_repeat_expr = None
		node_repeat_until = None
		node_size = None
		node_condition = None
		node_enum_type = None

		# Get values from node
		for key, value in child.items():
			if key == 'id':
				node_id = value
			if key == 'contents':
				if isinstance(value, str):
					node_contents = list(value.encode('ascii'))
				else:
					node_contents = value
			if key == 'type':
				if isinstance(value, str):
					node_type = get_type_identifier(value)
				else:
					for key_type, value_type in value.items():
						if key_type == 'switch-on':
							node_switch_on = prefix_expression(value_type)
						if key_type == 'cases':
							for case_enum, case_type in value_type.items():
								case_enum = str(case_enum)
								info = type_switch_info()
								node_type_switch.append(info)
								info.enumname = prefix_expression(case_enum)
								info.type = case_type
								if '::' in case_enum:
									split = case_enum.split('::')
									info.unionname = 'block_' + split[1]
								else:
									info.unionname = 'block_' + case_enum
								info.type = get_type_identifier(info.type)
			if key == 'repeat':
				node_repeat = value
			if key == 'repeat-expr':
				node_repeat_expr = prefix_expression(value)
			if key == 'repeat-until':
				node_repeat_until = prefix_expression(value)
			if key == 'size':
				node_size = prefix_expression(str(value))
			if key == 'if':
				node_condition = prefix_expression(value)
			if key == 'enum':
				node_enum_type = value

		# Generate code for loading sequence
		multi = ""
		substream = 's'
		struct_type = node_type;

		# Start condition block
		if node_condition is not None:
			self.add_code_all(f'if ({node_condition})')
			self.indent_code_plus_all()

		# Start repetition block
		if node_repeat is not None:
			multi='[i]'
			self.add_i()
			struct_type = f'struct {node_type}_ *'
			if node_repeat == 'expr':
				self.read.add_code(f'data->{node_id} = malloc(sizeof({node_type}) * {node_repeat_expr});')
				self.add_code_all(f'for (i = 0; i < {node_repeat_expr}; i++)')
				self.indent_code_plus_all()
			elif node_repeat == 'until' or node_repeat == 'eos':
				self.read.add_code(f'data->{node_id} = NULL;')
				self.add_code_all('i = -1;')
				self.add_code_all('do')
				self.indent_code_plus_all()
				self.add_code_all('i += 1;')
				self.read.add_code(f'data->{node_id} = realloc(data->{node_id}, sizeof({node_type}) * (i + 1));')
				self.add_struct(f'int {node_id}_count;')


		if node_enum_type is not None:
			struct_type = node_enum_type;


		if node_contents is not None:
			# Load fixed data that is always the same
			bytes = ', '.join(map(str, node_contents))
			self.add_var_readwrite(f'const char {node_id}[] = {{{bytes}}};')
			self.read.add_code(f'CHECK(stream_read_expect_bytes(s, {node_id}));')
			self.write.add_code(f'CHECK(stream_write_bytes(s, {node_id}, sizeof({node_id}), stage, do_sequence, _root));')
		elif node_size is not None and (node_type == None or node_type == 'str'):
			# Load dynamic, untyped data
			self.add_struct(f'char *{node_id};')
			self.read.add_code(f'data->{node_id} = malloc({node_size});')
			self.read.add_code(f'CHECK(stream_read_bytes(s, data->{node_id}, {node_size}));')
			self.write.add_code(f'CHECK(stream_write_bytes(s, data->{node_id}, {node_size}, stage, do_sequence, _root));')
			self.free.add_code(f'free(data->{node_id});')
		else:

			# Create substream ?
			if node_size is not None and node_type != None and node_type != 'str':
				self.add_var_readwrite(f'stream substream_{node_id};')
				self.read.add_code(f'CHECK(stream_read_make_substream(s, &substream_{node_id}, {node_size}));')
				self.write.add_code(f'CHECK(stream_write_make_substream(s, &substream_{node_id}));')
				substream = f'&substream_{node_id}'

			if node_switch_on is not None:
				# Load a type that is decided by an enum
				# Integer are for easier usage coerced into an int32
				is_int = node_type_switch[0].type.startswith('int')
				if is_int:
					self.add_struct(f'int32_t {node_id};')
				else:
					self.add_struct('union')
					self.indent_struct_plus()
				self.add_code_all(f'switch({node_switch_on})')
				self.indent_code_plus_all()
				for info in node_type_switch:
					if not is_int:
						self.add_struct(f'{info.type} {info.unionname};')
					self.add_dependency(info.type)
					self.add_code_all(f'case {info.enumname}:')
					self.indent_code_plus_all()
					if is_int:
						self.add_var_readwrite(f'{info.type} {node_id}_{info.unionname};')
						self.read.add_code(f'CHECK(stream_read_{info.type}(s_root, s, &{node_id}_{info.unionname}, _root));')
						self.read.add_code(f'data->{node_id} = {node_id}_{info.unionname};')
						self.write.add_code(f'{node_id}_{info.unionname} = data->{node_id};')
						self.write.add_code(f'CHECK(stream_write_{info.type}(s_root, s, &{node_id}_{info.unionname}, stage, do_sequence, _root));')
					else:
						self.read.add_code(f'CHECK(stream_read_{info.type}(s_root, s, &data->{info.unionname}));')
						self.write.add_code(f'CHECK(stream_write_{info.type}(s_root, s, &data->{info.unionname}, stage, do_sequence, _root));')
						self.free.add_code(f'stream_free_{info.type}(&data->{info.unionname}));')
					self.add_code_all('break;')
					self.indent_code_minus_all()
				self.indent_code_minus_all()
				if not is_int:
					self.indent_struct_minus_all('};')
			else:
				# Load a normal type
				self.add_struct(f'{struct_type} {node_id};')
				if node_repeat is None:
					self.add_dependency(node_type)
				if node_enum_type is not None:
					# For enums we need to make a copy, since not all enums are the same size (could read only an int8, but enum is int32)
					self.add_var_readwrite(f'{node_type} {node_id};')
					self.read.add_code(f'CHECK(stream_read_{node_type}(s_root, {substream}, &{node_id}, _root));')
					self.read.add_code(f'data->{node_id}{multi} = {node_id};')
					self.write.add_code(f'{node_id} = data->{node_id}{multi};')
					self.write.add_code(f'CHECK(stream_write_{node_type}(s_root, {substream}, &{node_id}, stage, do_sequence, _root));')
				else:
					self.read.add_code(f'CHECK(stream_read_{node_type}(s_root, {substream}, &data->{node_id}{multi}, _root));')
					self.write.add_code(f'CHECK(stream_write_{node_type}(s_root, {substream}, &data->{node_id}{multi}, stage, do_sequence, _root));')
					if not (node_type.startswith('int') or node_type.startswith('uint')):
						self.free.add_code(f'stream_free_{node_type}(&data->{node_id}{multi});')

			# Set substream size into struct
			if node_size is not None and node_type != None and node_type != 'str':
				expr = make_expression_for_setting_size(node_size, f'stream_write_get_length(&substream_{node_id})')
				if expr != None:
					self.write.add_code('if (stage == STREAM_WRITE_STAGE_DRYRUN && do_sequence)')
					self.write.indent_code_plus()
					self.write.add_code(expr)
					self.write.indent_code_minus()

		# Close repetition block
		if node_repeat == 'expr':
			self.indent_code_minus_all()
			self.free.add_code(f'free(data->{node_id});')
		elif node_repeat == 'until':
			self.indent_code_minus_all()
			repeat_cond = node_repeat_until.replace(f'##current##->', 'data->{node_id}[i].')
			self.add_code_all(f'while(!({repeat_cond}));')
			self.free.add_code(f'free(data->{node_id});')
		elif node_repeat == 'eos':
			self.indent_code_minus_all()
			self.read.add_code(f'while(s->start + s->pos < s->max);')
			self.read.add_code(f'data->{node_id}_count = i;')
			self.write.add_code(f'while(i + 1 < data->{node_id}_count);')
			self.free.add_code(f'while(i + 1 < data->{node_id}_count);')
			self.free.add_code(f'free(data->{node_id});')

		# Close condition block
		if node_condition is not None:
			self.indent_code_minus_all()

	def load_sequences(self, seq):
		self.write.add_code('/* Store position for current type */')
		self.write.add_code('if (stage == STREAM_WRITE_STAGE_DRYRUN && do_sequence)')
		self.write.indent_code_plus()
		self.add_struct('int _dryrun_pos;')
		self.write.add_code('data->_dryrun_pos = stream_write_get_position_absolute(s);')
		self.write.indent_code_minus()
		self.write.add_line()
		for child in seq:
			self.load_sequence(child)


def build_enums(node):
	ret = []
	for k in node:
		v = node[k]
		info = enum_info()
		ret.append(info)
		info.load(k, v)
	return ret

def build_types(node):
	ret = []
	for k in node:
		v = node[k]
		info = type_info(k)
		ret.append(info)
		info.load(v)
	return ret

def main():
	if len(sys.argv) != 2:
		print(f'Usage: {sys.argv[0]} file.ksy')
		return

	file_in = open(sys.argv[1], 'r')
	root = yaml.safe_load(file_in)
	file_in.close()

	meta_id = root['meta']['id']
	main_type = f'type_{meta_id}'

	enums = None
	mainSeq = seq_info()
	types = None
	


	type_main = type_info(main_type)

	# Need to reset the stream for write stage
	type_main.seq.write.add_code('if (stage == STREAM_WRITE_STAGE_WRITE)')
	type_main.seq.write.indent_code_plus()
	type_main.seq.write.add_code('s_root->pos = 0;');
	type_main.seq.write.indent_code_minus()

	type_main.load(root)

	for k in root:
		v = root[k]
		if k == 'types':
			types = build_types(v)
		if k == 'enums':
			enums = build_enums(v)
	types.append(type_main)

	types = sort_types(types)
	header = []
	source = []

	header.append('/* This file is generated - Do not edit manually */\n\n')
	header.append('#include "parser_generic.h"\n\n')
	if enums is not None:
		for enum in enums:
			header.append('typedef enum\n')
			header.append('{\n')
			for line in enum.lines_enum:
				header.append(line + '\n')
			header.append(f'}} {enum.name};\n\n')
	for type in types:
		write_struct(header, type.name, type.seq)
	header.append(f'int stream_read_{meta_id}(stream *s, {main_type} *data);\n')
	header.append(f'int stream_write_{meta_id}(stream *s, {main_type} *data, stream_write_stage stage);\n')
	header.append(f'void stream_free_{meta_id}({main_type} *data);\n')


	source.append('/* This file is generated - Do not edit manually */\n\n')
	source.append(f'#include "parser_{meta_id}.h"\n\n')
	for type in types:
		source.append(f'int stream_read_{type.name}(stream *s_root, stream *s, { type.name} *data, {main_type} *_root);\n')
		source.append(f'int stream_write_{type.name}(stream *s_root, stream *s, { type.name} *data, stream_write_stage stage, BOOL do_sequence, {main_type} *_root);\n')
		source.append(f'void stream_free_{type.name}({ type.name} *data);\n')
	source.append('\n')
	for type in types:
		write_method(source, type.name, type.seq.read, method_type.READ, main_type)
		write_method(source, type.name, type.seq.write, method_type.WRITE, main_type)
		write_method(source, type.name, type.seq.free, method_type.FREE, main_type)
	source.append(f'int stream_read_{meta_id}(stream *s, {main_type} *data)\n')
	source.append('{\n')
	source.append(f'\treturn stream_read_{main_type}(s, s, data, data);\n')
	source.append('}\n\n')
	source.append(f'int stream_write_{meta_id}(stream *s, { type.name} *data, stream_write_stage stage)\n')
	source.append('{\n')
	source.append(f'\tCHECK(stream_write_{main_type}(s, s, data, stage, TRUE, data));\n')
	source.append(f'\tCHECK(stream_write_{main_type}(s, s, data, stage, FALSE, data));\n')
	source.append('\treturn 0;\n')
	source.append('}\n\n')
	source.append(f'void stream_free_{meta_id}({ type.name} *data)\n')
	source.append('{\n')
	source.append(f'\tstream_free_{main_type}(data);\n')
	source.append('}\n')

	file_header = open(f'./tests/UIRibbon/parser_{meta_id}.h', 'w')
	file_header.write(''.join(header))
	file_header.close()
	file_source = open(f'./tests/UIRibbon/parser_{meta_id}.c', 'w')
	file_source.write(''.join(source))
	file_source.close()

main()