import argparse
import os
import re

from wprotocol_specs import *

parser = argparse.ArgumentParser()
parser.add_argument('dir', type=str,
                    help='Target directory')
args = parser.parse_args()
target_dir = args.dir


class MessageBuilder:
	def __init__(self):
		self.messages = {}
		self.current_message = None

	def comment(self, comment):
		None

	def new_message(self, name, line_data):
		if self.current_message != None:
			raise Exception("Attempt to process new message before completing last one!")
		self.current_message = {'name': name, 'line_data': line_data, 'fields': []}

	def field(self, fname, t, line_data):
		self.current_message['fields'].append((fname, t, line_data))

	def end_message(self):
		if self.current_message == None:
			raise Exception("Attempting to complete nonexistent message!")
		self.messages[self.current_message['name']] = self.current_message
		self.current_message = None

	def build_messages(self):
		return self.messages


class PonyLang:
	def __init__(self):
		self.spec = PonySpec()
		self.output = ''
		self.add_header()

	def add(self, strng, end='\n'):
		self.output += strng + end

	def compile(self):
		return self.output

	def add_messages(self, messages):
		self.spec.add_message_specs(messages)
		for (mname, message) in messages.items():
			fields = message['fields']
			self.add_class(mname, fields)
			self.add_encoder(mname, fields)
			self.add_decoder(mname, fields)
			self.add_string_output_encoder(mname)
		self.add('')

	def add_header(self):
		self.add('\
""" \n\
*******************************\n\
THIS IS GENERATED CODE.\n\
DO NOT EDIT THIS FILE DIRECTLY.\n\
EDIT .wproto FILES INSTEAD.\n\
*******************************\n\
"""\n')
		self.add('use "buffered"')
		self.add('use "collections"\n')

	def add_class(self, mname, fields):
		self.add_name(mname)
		for (fname, t, line_data) in fields:
			self.add_field(fname, t)
		self.add('')
		self.add_class_methods(mname, fields)

	def add_name(self, n):
		self.add("class val " + n + " is (Hashable & Equatable[" + n + "])")

	def add_field(self, fname, t):
		self.add("  let {}: {}".format(fname, self.spec.pony_type_for(t)))

	def add_class_methods(self, mname, fields):
		self.add_constructor(fields)
		self.add_eq(mname, fields)
		self.add_hash(fields)
		self.add_tostring(mname, fields)

	def add_constructor(self, fields):
		prime_parameters = ", ".join(map(lambda pair: pair[0] + "': " + self.spec.pony_type_for(pair[1]), fields))
		self.add("  new val create(" + prime_parameters + ") =>")
		for (fname, t, _) in fields:
			self.add("    " + fname + " = " + fname + "'")
		self.add("")

	def add_eq(self, mname, fields):
		self.add("  fun eq(that: box->" + mname + "): Bool =>")
		equals = " and ".join(map(lambda pair: self.spec.eq(pair[0], pair[1]), fields))
		self.add("    " + equals)
		self.add("")

	def add_hash(self, fields):
		self.add("  fun hash(): USize =>")
		hashes = " xor ".join(map(lambda pair: self.spec.hash(pair[0], pair[1]), fields))
		self.add("    (" + hashes + ").usize()")
		self.add("")

	def add_tostring(self, mname, fields):
		self.add("  fun string(): String iso^ =>")
		stringy = ' + "|'.join(map(lambda pair: pair[0] + ':" + ' + self.spec.tostring(pair[0], pair[1]), fields))
		self.add('    ("' + mname + '(' + stringy + ' + ")").clone()')
		self.add("")

	def add_encoder(self, mname, fields):
		self.add("primitive " + mname + "Encoder")
		self.add("  fun apply(" + mname.lower() + ": " + mname + ", wb: Writer = Writer): Array[ByteSeq] val =>")
		self_fields = ", ".join(map(lambda pair: mname.lower() + "." + pair[0], fields))
		self.add("    from_fields(wb, " + self_fields + ")")
		self.add("")
		parameters = ", ".join(map(lambda pair: pair[0] + ": " + self.spec.pony_type_for(pair[1]), fields))
		self.add("  fun from_fields(wb: Writer, " + parameters + "): Array[ByteSeq] val =>")
		self.add("    let payload_size: USize = " + self.spec.size('', mname)) #" + ".join(self.payload_size))
		self.add("    wb.u32_be(payload_size.u32())")
		self.add("    var next_size: USize = 0")
		for (fname, t, _) in fields:
			self.add(self.spec.encode(fname, t), end="")
		self.add("    wb.done()")
		self.add("")

	def add_decoder(self, mname, fields):
		self.add("primitive " + mname + "Decoder")
		self.add("  fun apply(data: Array[U8] val, rb: Reader = Reader): " + mname + " ? =>")
		self.add("    var next_size: USize = 0")
		self.add("    rb.append(data)")
		for (fname, t, _) in fields:
			self.add(self.spec.decode(fname, t), end="")
		args = ", ".join(map(lambda pair: pair[0], fields))
		self.add("    " + mname + "(" + args + ")")
		self.add("")

	def add_string_output_encoder(self, mname):
		self.add("primitive " + mname + "StringOutputEncoder")
		self.add("  fun apply(" + mname.lower() + ": " + mname + ", wb: Writer = Writer): Array[ByteSeq] val =>")
		self.add("    wb.write(" + mname.lower() + '.string() + "\\n")')
		self.add("    wb.done()")
		self.add("")

blankp = re.compile(r'^s*$')
commentp = re.compile(r'^\s*#(.*)$')
namep = re.compile(r'^(message)\s*([A-Z_][A-Za-z0-9_]+)\s*{\s*$') # start message mode
badnamep = re.compile(r'^(message)\s*([a-z0-9]\S*)\s*{\s*$') # start message mode
badname2p = re.compile(r'^(message)\s*([A-Z_]\S*[^\w]\S*)\s*{\s*$') # start message mode
fieldp = re.compile(r'^\s*([a-z_]+[A-Za-z0-9_]*):\s*([A-Z_][A-Za-z0-9_\[\]]+)\s*$') # in message mode: fieldtype field
endp = re.compile(r'^\s*}\s*$') # <--- stop message mode
annotationp = re.compile(r'^\s*@([A-Za-z0-9_]+)\s*$')

def raise_syntax_error(filename, line_no, line, msg):
	raise Exception("Syntax error, " + filename + ":" + str(line_no) + "!\n\n" + line + "\n" + msg + "\n")

def generate_code_for(wproto_file, lang):
	line_no = 0
	message_mode = False
	msg_builder = MessageBuilder()
	for line in wproto_file:
		line_no += 1
		blank_res = re.match(blankp, line)
		comment_res = re.match(commentp, line)
		name_res = re.match(namep, line)
		bad_name_res = re.match(badnamep, line)
		bad_name2_res = re.match(badname2p, line)
		field_res = re.match(fieldp, line)
		annotation_res = re.match(annotationp, line)
		end_res = re.match(endp, line)

		if bool(blank_res):
			None
		elif bool(comment_res):
			None
		elif bool(name_res):
			if message_mode:
				raise_syntax_error(wproto_file.name, line_no, line, "Inside another message definition!")
			message_mode = True
			name = name_res.group(2)
			msg_builder.new_message(name, (line_no, line))
		elif bool(bad_name_res):
			raise_syntax_error(wproto_file.name, line_no, line, "Message name must begin with A-Z or _!")
		elif bool(bad_name2_res):
			raise_syntax_error(wproto_file.name, line_no, line, "Message name must only contain A-Z, a-z, or _!")
		elif bool(field_res):
			if not message_mode:
				raise_syntax_error(wproto_file.name, line_no, line, "Outside message definition!")
			msg_builder.field(field_res.group(1), field_res.group(2), (line_no, line))
		elif bool(end_res):
			if not message_mode:
				raise_syntax_error(wproto_file.name, line_no, line, "Outside message definition!")
			message_mode = False
			msg_builder.end_message()
		elif bool(annotation_res):
			None
		else:
			raise_syntax_error(wproto_file.name, line_no, line, "Invalid line!")

	lang.add_messages(msg_builder.build_messages())


# START
lang = PonyLang()
for filename in os.listdir(target_dir):
    if filename.endswith(".wproto"):
    	with open(target_dir + "/" + filename, "r") as wproto_file:
    		generate_code_for(wproto_file, lang)

compiled = lang.compile()
print(compiled)
