# 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 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 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

	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, filename_hi: str, filename_lo: str, out_dir="."):
		self.out_dir = out_dir
		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)


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

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

	@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

	@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



class PpgEvu1v00(AImage):
	pass

class PpgEvu2v00(AImage):
	pass


def detect(args):
	os.makedirs(args.output_dir, exist_ok=True)
	image = WaldorfMicroWave2V00(args.filepath, args.extra, args.output_dir)
	pass

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


if __name__ == '__main__':
	parser = ArgumentParser()
	parser.add_argument("filepath", help="Location of firmware binary image")
	parser.add_argument("--extra", help="Location of 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/