#!/usr/bin/env python3

from collections import defaultdict

from datetime import datetime

from gzip import GzipFile

from sys import stderr

from xml.etree.ElementTree import SubElement, parse as parseXML

import platform

cpuClock = 3579545

masterClock = cpuClock * 960

def readStates(topFilename):

	openFilenames = []

	base = None

	out = None

	scale = 1

	inputMap = {}

	inputStates = []

	time = 0

	def handleProcessingInstruction(words):

		if not words:

			raise ValueError('Missing processing instruction')

		elif words[0] == 'base':

			if len(words) != 2:

				raise ValueError(

						'Processing instruction "base" expects 1 argument, '

						'got %d' % (len(words) - 1))

			nonlocal base

			if base is not None:

				raise ValueError('Attempt to change base file name')

			base = words[1]

		elif words[0] == 'out':

			if len(words) != 2:

				raise ValueError(

						'Processing instruction "out" expects 1 argument, '

						'got %d' % (len(words) - 1))

			nonlocal out

			if out is not None:

				raise ValueError('Attempt to change output file name')

			out = words[1]

		elif words[0] == 'scale':

			if len(words) != 2:

				raise ValueError(

						'Processing instruction "scale" expects 1 argument, '

						'got %d' % (len(words) - 1))

			nonlocal scale

			scale = int(words[1])

			print('time scale: %.2f frames per second' % (masterClock / scale))

		elif words[0] == 'input':

			if len(words) != 5:

				raise ValueError(

						'Processing instruction "input" expects 4 arguments, '

						'got %d' % (len(words) - 1))

			name, inpType, rowStr, colStr = words[1:]

			if inpType != 'key':

				raise ValueError('Unknown input type "%s"' % inpType)

			row = int(rowStr)

			col = int(colStr)

			inputMap[name] = (row, col)

		elif words[0] == 'include':

			if len(words) != 2:

				raise ValueError(

						'Processing instruction "include" expects 1 argument, '

						'got %d' % (len(words) - 1))

			readFile(words[1])

		else:

			raise ValueError('Unknown processing instruction: %s' % words[0])

	def handleState(words):

		frames = int(words[0])

		try:

			inputs = [inputMap[name] for name in words[1:]]

		except KeyError as ex:

			raise ValueError('Undefined input: %s' % ex)

		nonlocal time

		inputStates.append((time, inputs))

		time += frames * scale

		return frames

	def readFile(filename):

		if filename in openFilenames:

			raise ValueError('Circular include: %s' % filename)

		localStateCount = 0

		localFrameCount = 0

		startTime = time

		openFilenames.append(filename)

		try:

			with open(filename, 'r') as stream:

				for lineNr, line in enumerate(stream, 1):

					try:

						line = line.strip()

						if not line or line[0] == '#':

							# Ignore empty line or comment.

							continue

						if line[0] == '=':

							handleProcessingInstruction(line[1:].split())

						else:

							localFrameCount += handleState(line.split())

							localStateCount += 1

					except ValueError:

						role = 'In' if filename == openFilenames[-1] \

								else 'included from'

						print('%s "%s" line %d,' % (role, filename, lineNr),

							file=stderr)

						raise

		except OSError as ex:

			print('Failed to open input file "%s",' % filename, file=stderr)

			raise ValueError(str(ex)) from ex

		del openFilenames[-1]

		localTime = time - startTime

		print('%-17s %5d states, %6d frames, %7.2f seconds' % (

				filename + ':', localStateCount, localFrameCount,

				localTime / masterClock

				), file=stderr)

	readFile(topFilename)

	inputStates.append((time, []))

	if base is None:

		raise ValueError('Base file not defined')

	return base, out, inputStates

def statesToEvents(inputStates):

	active = {}

	for time, state in inputStates:

		stateByRow = defaultdict(int)

		for row, col in state:

			stateByRow[row] |= 1 << col

		for row in sorted(active.keys() | stateByRow.keys()):

			old = active.get(row, 0)

			new = stateByRow.get(row, 0)

			press = new & ~old

			release = ~new & old

			if press != 0 or release != 0:

				yield time, row, press, release

			active[row] = new

def replaceEvents(inp, out, inputEvents):

	doc = parseXML(inp)

	# Set the serialization date to now.

	rootElem = doc.getroot()

	rootElem.attrib['date_time'] = \

			datetime.now().strftime('%a %b %d %H:%M:%S %Y')

	rootElem.attrib['openmsx_version'] = 'txt2omr'

	rootElem.attrib['platform'] = platform.system().lower()

	# Remove snapshots except the one at timestamp 0.

	snapshots = doc.find('replay/snapshots')

	if snapshots is None:

		print('Base replay lacks snapshots', file=stderr)

	else:

		seenInitialSnapshot = False

		for snapshot in snapshots.findall('item'):

			timeElem = snapshot.find('scheduler/currentTime/time')

			time = int(timeElem.text)

			if time == 0:

				seenInitialSnapshot = True

			else:

				snapshots.remove(snapshot)

		if not seenInitialSnapshot:

			print('No snapshot found with timestamp 0', file=stderr)

	# Replace event log.

	eventsElem = doc.find('replay/events')

	if eventsElem is None:

		print('No events tag found; cannot insert events', file=stderr)

	else:

		tail = eventsElem.tail

		eventsElem.clear()

		eventsElem.text = '\n'

		eventsElem.tail = tail

		# IDs must be unique for the entire document. We look for the highest

		# in-use ID and generate new IDs counting up from there.

		baseID = max(

			(

				int(id_val)

				for elem in doc.iterfind('.//*[@id]')

				if (id_val := elem.attrib['id']).isdecimal()

			),

			default=-1

		) + 1

		def createEvent(i, time):

			itemElem = SubElement(eventsElem, 'item',

					id=str(baseID + i), type='KeyMatrixState')

			itemElem.tail = '\n'

			stateChangeElem = SubElement(itemElem, 'StateChange')

			timeElem1 = SubElement(stateChangeElem, 'time')

			timeElem2 = SubElement(timeElem1, 'time')

			timeElem2.text = str(time)

			return itemElem

		for i, (time, row, press, release) in enumerate(inputEvents):

			itemElem = createEvent(i, time)

			SubElement(itemElem, 'row').text = str(row)

			SubElement(itemElem, 'press').text = str(press)

			SubElement(itemElem, 'release').text = str(release)

		endTime = inputEvents[-1][0] if inputEvents else 0

		createEvent(len(inputEvents), endTime).attrib['type'] = 'EndLog'

	# Reset re-record count.

	reRecordCount = doc.find('replay/reRecordCount')

	if reRecordCount is not None:

		reRecordCount.text = '0'

	# Reset the current time.

	currentTime = doc.find('replay/currentTime/time')

	if currentTime is not None:

		currentTime.text = '0'

	out.write(b'<?xml version="1.0" ?>\n')

	out.write(b"<!DOCTYPE openmsx-serialize SYSTEM 'openmsx-serialize.dtd'>\n")

	doc.write(out, encoding='utf-8', xml_declaration=False)

def convert(inFilename):

	try:

		base, outFilename, inputStates = readStates(inFilename)

	except ValueError as ex:

		print('ERROR: %s' % ex, file=stderr)

		exit(1)

	inputEvents = list(statesToEvents(inputStates))

	def createOutput(inp):

		if outFilename is None:

			print('No output file name set', file=stderr)

		else:

			try:

				with GzipFile(outFilename, 'wb') as out:

					replaceEvents(inp, out, inputEvents)

			except OSError as ex:

				print('Failed to open output replay:', ex)

				exit(1)

			else:

				print('wrote output replay:', outFilename)

	try:

		if base.endswith('.omr'):

			with GzipFile(base, 'rb') as inp:

				createOutput(inp)

		elif base.endswith('.xml'):

			with open(base, 'rb') as inp:

				createOutput(inp)

		else:

			print('Unknown base file type in "%s" '

					'(".xml" and ".omr" are supported)'

					% base, file=stderr)

			exit(1)

	except OSError as ex:

		print('Failed to open base replay:', ex, file=stderr)

		exit(1)

if __name__ == '__main__':

	from sys import argv

	if len(argv) != 2:

		print('Usage: txt2omr.py replay.txt', file=stderr)

		print('Converts the text version of an openMSX replay to an OMR file.',

				file=stderr)

		exit(2)

	else:

		convert(argv[1])