"""Build the verified V21 texture-reclaim runtime DLL from a local game DLL."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
import hashlib
from pathlib import Path
import shutil
import struct


SOURCE_SHA256 = "4b4e5c28e2bfdaa8d36c8341fad27becfd03faf63aa8ea874f5d43c5b1961868"
PATCHED_SHA256 = "0be11b06ccf9be747e32b195c8d5980062f446d7d3a0a662ebfe0e602a01c538"
TEXT_VA = 0x10001000
TEXT_FILE_OFFSET = 0x400
ENTRY_VA = 0x107C0B40
ENTRY_FILE_OFFSET = TEXT_FILE_OFFSET + (ENTRY_VA - TEXT_VA)
ENTRY_ORIGINAL = bytes.fromhex("55 8b ec 6a ff")
CAVE_VA = 0x11195CE0
CAVE_FILE_OFFSET = TEXT_FILE_OFFSET + (CAVE_VA - TEXT_VA)
TEXT_RAW_SIZE = 0x1194E00
MAX_HOOK_SIZE = TEXT_FILE_OFFSET + TEXT_RAW_SIZE - CAVE_FILE_OFFSET
HOOK = bytes.fromhex(
	"9c60e8490136ff85c0744a89c38d70088dbe505d010039fe743b80be58030000"
	"00742a8b068b4c242885c974083b81d41f000074188b4c242c85c974083b81d4"
	"1f000074085089d9e8a32036ff81c650040000ebc1619d5589e56affe904ae62"
	"ff"
)


class TextureReclaimPatchError(RuntimeError):
	"""Raised when the exact V21 binary profile cannot be patched safely."""


@dataclass(frozen=True)
class TextureReclaimPatchResult:
	source: str
	destination: str
	source_sha256: str
	destination_sha256: str
	entry_va: str
	cave_va: str
	hook_size: int
	old_text_virtual_size: int
	new_text_virtual_size: int
	commands: tuple[tuple[str, ...], ...]
	created_at: str


def build_texture_reclaim_patch(
	source: Path,
	destination: Path,
) -> TextureReclaimPatchResult:
	"""Create the V21 patch in a separate runtime cache, never in the game tree."""
	source = source.expanduser().resolve(strict=True)
	if _sha256(source) != SOURCE_SHA256:
		raise TextureReclaimPatchError(
			"source game.dll does not match the verified V21 profile"
		)

	hook = HOOK
	commands: list[tuple[str, ...]] = []
	if not hook or len(hook) > MAX_HOOK_SIZE:
		raise TextureReclaimPatchError(
			f"compiled hook size {len(hook)} exceeds the {MAX_HOOK_SIZE}-byte code cave"
		)

	image = bytearray(source.read_bytes())
	if bytes(image[ENTRY_FILE_OFFSET : ENTRY_FILE_OFFSET + len(ENTRY_ORIGINAL)]) != ENTRY_ORIGINAL:
		raise TextureReclaimPatchError(
			"V21 map-loader entry bytes do not match the verified profile"
		)
	if any(image[CAVE_FILE_OFFSET : CAVE_FILE_OFFSET + len(hook)]):
		raise TextureReclaimPatchError("V21 code cave is not empty")

	text_header = _find_pe_section(image, b".text")
	old_virtual_size, text_rva, raw_size, raw_offset = struct.unpack_from(
		"<IIII", image, text_header + 8
	)
	if (text_rva, raw_size, raw_offset) != (
		TEXT_VA - 0x10000000,
		TEXT_RAW_SIZE,
		TEXT_FILE_OFFSET,
	):
		raise TextureReclaimPatchError(
			"V21 executable section layout does not match the verified profile"
		)
	new_virtual_size = CAVE_VA - TEXT_VA + len(hook)
	if new_virtual_size > TEXT_RAW_SIZE:
		raise TextureReclaimPatchError("compiled hook extends beyond the executable section")
	if old_virtual_size > new_virtual_size:
		raise TextureReclaimPatchError(
			"verified executable section is already larger than the patch target"
		)

	entry_displacement = CAVE_VA - (ENTRY_VA + 5)
	image[ENTRY_FILE_OFFSET : ENTRY_FILE_OFFSET + 5] = b"\xe9" + struct.pack(
		"<i", entry_displacement
	)
	image[CAVE_FILE_OFFSET : CAVE_FILE_OFFSET + len(hook)] = hook
	struct.pack_into("<I", image, text_header + 8, new_virtual_size)

	destination = destination.expanduser().absolute()
	destination.parent.mkdir(parents=True, exist_ok=True)
	temporary = destination.with_name(f".{destination.name}.tmp")
	try:
		temporary.write_bytes(image)
		if _sha256(temporary) != PATCHED_SHA256:
			raise TextureReclaimPatchError(
				"built runtime patch does not match the verified V21 profile"
			)
		temporary.replace(destination)
		shutil.copystat(source, destination)
	except BaseException:
		temporary.unlink(missing_ok=True)
		raise

	return TextureReclaimPatchResult(
		source=str(source),
		destination=str(destination),
		source_sha256=SOURCE_SHA256,
		destination_sha256=PATCHED_SHA256,
		entry_va=f"0x{ENTRY_VA:08x}",
		cave_va=f"0x{CAVE_VA:08x}",
		hook_size=len(hook),
		old_text_virtual_size=old_virtual_size,
		new_text_virtual_size=new_virtual_size,
		commands=tuple(commands),
		created_at=datetime.now().astimezone().isoformat(),
	)


def _find_pe_section(image: bytes | bytearray, section_name: bytes) -> int:
	if len(image) < 0x40:
		raise TextureReclaimPatchError("source is too small to be a PE image")
	pe_offset = struct.unpack_from("<I", image, 0x3C)[0]
	if image[pe_offset : pe_offset + 4] != b"PE\0\0":
		raise TextureReclaimPatchError("source does not have a valid PE signature")
	number_of_sections = struct.unpack_from("<H", image, pe_offset + 6)[0]
	optional_header_size = struct.unpack_from("<H", image, pe_offset + 20)[0]
	section_table = pe_offset + 24 + optional_header_size
	for index in range(number_of_sections):
		header = section_table + index * 40
		name = bytes(image[header : header + 8]).rstrip(b"\0")
		if name == section_name:
			return header
	raise TextureReclaimPatchError(f"PE section is missing: {section_name!r}")


def _sha256(path: Path) -> str:
	digest = hashlib.sha256()
	with path.open("rb") as handle:
		for chunk in iter(lambda: handle.read(1024 * 1024), b""):
			digest.update(chunk)
	return digest.hexdigest()