"""
Handle ZynqMP boot images, as generated by the Xilinx bootgen tool.
"""

import logging

logger = logging.getLogger("snagrecover")
from snagrecover.protocols import dfu
from snagrecover.utils import BinFileHeader as Header
from dataclasses import dataclass, astuple
import struct


@dataclass
class ZynqMPImageTable(Header):
	version: int
	part_count: int
	part_hdr: int
	image_hdr: int
	auth_cert: int
	aux_bootdev: int
	padding: bytes
	csum: int

	fmt = "<LLLLLL36sL"
	class_size = 0x40

	def update_checksum(self):
		data = struct.pack(ZynqMPImageTable.fmt, *astuple(self))
		self.csum = zynqmp_csum(data[:-4])


@dataclass
class ZynqMPImageHeader(Header):
	next_image: int
	part_hdr: int
	reserved: int
	part_count: int
	# Other fields are ignored

	fmt = "<LLLL"
	class_size = 16


@dataclass
class ZynqMPPartHeader(Header):
	enc_size: int
	plain_size: int
	total_size: int
	next_part: int
	exec_lo: int
	exec_hi: int
	load_lo: int
	load_hi: int
	start: int
	attrs: int
	sec_count: int
	csum_word_offset: int
	image_hdr: int
	ac_offset: int
	part_id: int
	csum: int

	fmt = "<LLLLLLLLLLLLLLLL"
	class_size = 0x40

	def update_checksum(self):
		data = struct.pack(ZynqMPPartHeader.fmt, *astuple(self))
		self.csum = zynqmp_csum(data[:-4])


def zynqmp_csum(data):
	s = 0

	for i in range(0, len(data), 4):
		s = (s + int.from_bytes(data[i : i + 4], "little")) % (1 << 32)

	return s ^ 0xFFFFFFFF


def find_img_table(boot_bin):
	boot_hdr = boot_bin[:0xB5]
	hdr = boot_hdr

	sig = hdr[0x24:0x28]
	if sig != b"XNLX":
		raise ValueError("Invalid ZynqMP boot header signature!")

	return int.from_bytes(boot_hdr[0x98:0x9C], "little")


def drop_images(boot_bin: bytearray, keep_images):
	"""
	Drop partitions which aren't required for the FSBL stage from the boot
	image.
	"""

	if keep_images == 0:
		raise ValueError("Must keep at least one image!")

	img_table_offset = find_img_table(boot_bin)
	img_table = ZynqMPImageTable.read(boot_bin, img_table_offset)

	kept_images = []
	next_image = 4 * img_table.image_hdr
	for _ in range(keep_images):
		image = ZynqMPImageHeader.read(boot_bin, next_image)
		kept_images.append(image)
		next_image = 4 * image.next_image

	last_kept_image = kept_images[-1]
	if next_image == 0:
		return

	# Mark as last image
	last_kept_image.next_image = 0
	ZynqMPImageHeader.write(last_kept_image, boot_bin)

	# Find last remaining partition
	next_part = 4 * last_kept_image.part_hdr
	for _ in range(last_kept_image.part_count):
		part = ZynqMPPartHeader.read(boot_bin, next_part)
		next_part = 4 * part.next_part

	last_kept_part = part

	# Drop unused partition data
	cutoff = 4 * (last_kept_part.start + last_kept_part.total_size)

	# Mark as last partition and update csum
	last_kept_part.next_part = 0
	last_kept_part.update_checksum()
	ZynqMPPartHeader.write(last_kept_part, boot_bin)

	# Update partition count and csum
	img_table.part_count = sum([image.part_count for image in kept_images])
	img_table.update_checksum()

	ZynqMPImageTable.write(img_table, boot_bin)

	return bytes(boot_bin[:cutoff])


def zynqmp_run(dev, fw_name, fw_blob, subfw_name):
	dfu_cmd = dfu.DFU(dev, stm32=False)

	if fw_name == "fsbl":
		partid = 0
	if fw_name == "boot":
		if subfw_name == "fsbl":
			partid = 0
			"""
			Only the FSBL and PMUFW are required for the first stage.
			Keep the first two images and drop the rest, so that
			the firmware blob can fit in SRAM.
			"""
			fw_blob = drop_images(bytearray(fw_blob), 2)
		else:
			partid = 1

	logger.info("Downloading file...")
	dfu_cmd.download_and_run(fw_blob, partid, offset=0, size=len(fw_blob))
	logger.info("Done")
