# This is a sample Python script.
import enum
# Press Umschalt+F10 to execute it or replace it with your code.
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.

import os
import re
import struct
from argparse import ArgumentParser
from io import BytesIO
from typing import BinaryIO
import hashlib
import json

from pygccxml.declarations import unsigned_int_t


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

	def parse(self):
		pass

	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)
		filename_hi = None
		filename_lo = None
		for filename in os.listdir(folder):
			filepath = os.path.join(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

		if filename_lo is None:
			return

		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)


	@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)
		filename_hi = None
		filename_lo = None
		for filename in os.listdir(folder):
			filepath = os.path.join(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

		if filename_lo is None:
			return

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


def detect(args):
	image = WaldorfMicroWave2V00(args.folder, args.output_dir)
	image = WaldorfMicrowaveXT(args.folder, args.output_dir)
	pass

def main(args):
	image_type = detect(args)


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/