import enum
import os
import re
import struct
from argparse import ArgumentParser
import hashlib
import json
import abc

class AImage(abc.ABC):
	def __init__(self, folder: str, output_dir: str = "out"):
		self.folder = folder
		self.out_path = os.path.join(folder, output_dir)
		os.makedirs(self.out_path, exist_ok=True)

	def parse(self):
		return False

	def version(self):
		pass

	def is_valid(self):
		return False

	@staticmethod
	def _check(data: bytes, target: dict):
		is_valid = True
		is_valid &= AImage._check_len(data, target['LEN'])
		is_valid &= AImage._check_fingerprint(data, target['DIGEST'])
		is_valid &= AImage._check_version(data, 0x1B430, bytearray.fromhex(target['VERSION']))

		print("is_valid =", is_valid)

		return is_valid

	@staticmethod
	def _check_len(data: bytes, target: int):
		return len(data) == target

	@staticmethod
	def _check_fingerprint(data: bytes, target: str):
		h = hashlib.sha256()
		h.update(data)
		actual = h.digest()
		print(actual)
		return actual == target

	@staticmethod
	def _check_version(data: bytes, offset : int, target: bytearray):
		actual = data[offset:offset+len(target)]
		print(actual)
		return actual == target

	@staticmethod
	def combine(filename_hi: str, filename_lo: str):
		data_hi = b''
		data_lo = b''
		with open(filename_hi, "rb") as fp:
			data_hi = fp.read()
		with open(filename_lo, "rb") as fp:
			data_lo = fp.read()

		data = b''
		for hi, lo in zip(data_hi, data_lo):
			z = int.from_bytes([hi, lo], byteorder='big')
			b = struct.pack('>H', z)
			data += b

		return data, data_hi, data_lo

	def _save(self, filename: str, data):
		with open(os.path.join(self.out_path, filename), 'wb') as fp:
			fp.write(data)

	def _json_save(self, filename: str, data):
		with open(os.path.join(self.out_path, filename), 'w') as fp:
			json.dump(data, fp, indent=4)

	def get_waves(self):
		pass

	def get_wave_tables(self):
		pass

class WaldorfMicroWave1V20(AImage):
	pass


class WaldorfMicroWave1V25(AImage):
	pass

class WaldorfMicroWave2V00(AImage):
	class ImagePart(enum.Enum):
		Hi = 0
		Lo = 1

	target = {
		"LEN" : 131072,
		"DIGEST" : b'\xe3G\n~\xfb\xd0P\xf4~\x8d\xf3\x12p9\x88X\x95e\xc2\xc5e\xdev\xe6\r\xaba\x8b\xadDY\x90',
		"VERSION" : "322E3030300000004D6F6E204D617220372031323A30363A3132204D4554203139393400303230303934303337000000"
	}

	wave_data_descr = [
		{
			'offset': 0xAA18,
			'len': 0x1B00
		},
		{
			'offset': 0xAA18,
			'len': 0x3000
		}
	]

	wave_tables_descr = [
		{
			'offset': 0x9A18,
			'len': 0x0E00
		},
		{
			'offset': 0x9A18,
			'len': 0x0600
		}
	]

	def __init__(self, folder: str, out_dir="."):
		AImage.__init__(self, folder, out_dir)
		self._is_valid = False

	def parse(self):
		filename_hi = None
		filename_lo = None
		for filename in os.listdir(self.folder):
			filepath = os.path.join(self.folder, filename)
			if re.match('.*_H.bin', filename):
				filename_hi = filepath
			if re.match('.*_L.bin', filename):
				filename_lo = filepath

		if filename_hi is None:
			return False

		if filename_lo is None:
			return False

		data, data_hi, data_lo = self.combine(filename_hi, filename_lo)

		self._is_valid = self._check(data, self.target)
		if self._is_valid:
			self._save("combined_test_be.bin", data)

			wave_data = self.wave_extract(self.wave_data_descr, data_hi, data_lo)
			self._save("wave_data.bin", wave_data)

			wave_table = self.wave_table_extract(self.wave_tables_descr, data_hi, data_lo)
			self._json_save("wave_table.json", wave_table)

		return self._is_valid

	@staticmethod
	def wave_extract(descr, data_hi, data_lo):
		data = b''
		for des, rom_data in zip(descr, [data_hi, data_lo]):
			offset = des['offset']
			length = des['len']
			data += rom_data[offset:offset+length]
		return data

	@staticmethod
	def wave_table_extract(descr, data_hi, data_lo):
		data = []
		table_count = 0
		pos = 0
		for des, rom_data in zip(descr, [data_hi, data_lo]):
			offset = des['offset']
			length = des['len']

			for addr in range(offset, offset + length, 2):
				index = struct.unpack('>h', rom_data[addr:addr+2])[0]
				if index > 0:
					if pos == 0:
						print(f"----------------------------")
						print(f"- table #{table_count:03d}")
						print(f"----------------------------")
						table_count += 1
					entry = {"index": index, "pos": pos}
					print(entry)
					data.append(entry)

				if pos < 63:
					pos += 1
				else:
					pos = 0
		return data

	def is_valid(self):
		return self._is_valid


class PpgEvu1v00(AImage):
	pass

class PpgEvu2v00(AImage):
	pass

class WaldorfMicrowaveXT(AImage):
	def __init__(self, folder: str, out_dir="."):
		AImage.__init__(self, folder, out_dir)

	def parse(self):
		filename_hi = None
		filename_lo = None
		for filename in os.listdir(self.folder):
			filepath = os.path.join(self.folder, filename)
			if re.match('upper_.*.bin', filename):
				filename_hi = filepath
			if re.match('lower_.*.bin', filename):
				filename_lo = filepath

		if filename_hi is None:
			return False

		if filename_lo is None:
			return False

		data, data_hi, data_lo = self.combine(filename_lo, filename_hi)
		self._save("combined_be.bin", data)

		return True

def detect(args):
	wave_obj_list = [
		WaldorfMicroWave2V00(args.folder, args.output_dir),
		WaldorfMicrowaveXT(args.folder, args.output_dir)
	]

	for wave_obj in wave_obj_list:
		if wave_obj.parse():
			return wave_obj

	return None

def main(args):
	result = detect(args)
	print(f"Detected parser: {result}")


if __name__ == '__main__':
	parser = ArgumentParser()
	parser.add_argument("folder", help="Folder with firmware binary image")
	parser.add_argument("--output_dir", help="Location of output data", default='out')

	args = parser.parse_args()

	main(args)



# See PyCharm help at https://www.jetbrains.com/help/pycharm/