#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2026 Tomasz Kramkowski <tomasz@kramkow.ski>
# SPDX-License-Identifier: CC0-1.0

"""A Half-Life: Alyx launch wrapper that patches the game.

For build 26283546.

Use at your own risk!
"""

import hashlib
import os
import stat
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path

NAME = Path(__file__).name
MARKER = b"Valve, please fix."


@dataclass(frozen=True)
class Edit:
    offset: int
    replacement: bytes


@dataclass(frozen=True)
class Patch:
    original_hash: bytes
    patched_hash: bytes
    edits: tuple[Edit, ...]

    def apply_to(self, data: bytes) -> bytes:
        payload = data.removesuffix(MARKER)
        digest = sha256(payload)
        if digest == self.patched_hash:
            return payload + MARKER
        if digest != self.original_hash:
            if payload != data:
                raise LookupError(
                    "already patched with another set of patches: verify file integrity and restart"
                )
            raise LookupError("unrecognized build: this wrapper may need updating")
        result = bytearray(payload)
        for edit in self.edits:
            end = edit.offset + len(edit.replacement)
            if edit.offset < 0 or end > len(result):
                raise ValueError("patch extends outside the file")
            result[edit.offset : end] = edit.replacement
        if sha256(result) != self.patched_hash:
            raise ValueError("patched file checksum mismatch")
        return bytes(result) + MARKER


def patch(
    original_hash: str, patched_hash: str, edits: tuple[tuple[int, str], ...]
) -> Patch:
    return Patch(
        bytes.fromhex(original_hash),
        bytes.fromhex(patched_hash),
        tuple(
            Edit(offset, bytes.fromhex(replacement)) for offset, replacement in edits
        ),
    )


PATCHES = {
    "game/bin/linuxsteamrt64/librendersystemvulkan.so": patch(
        "dc034f6b5d2116393bfead77f3265c124687b76766906c53a077c196b21ffa18",
        "7a52bda8f4e6cc35e5ed46f181c3847400807704cab3fea38b40154d0b977f03",
        (
            # The renderer keeps a pool of "staging" buffers that it uses to
            # upload data to the GPU. When this pool exceeds 128MiB
            # (configurable via RenderSystem.VulkanStagingPMBSizeLimitMB in
            # "gameinfo.gi") of allocations, it will clear it down almost
            # completely. This can take a good 30+ms which makes for a pretty
            # unacceptable stutter. There's no need to clear all of it at once.
            # This patch stops clearing when one staging buffer has been freed.
            (0x22545D, "eb7b90"),
            # When trying to re-use a staging buffer, the renderer only looks
            # at two candidate buffers, and allocates (potentially triggering
            # the de-allocation above) otherwise. Searching all the buffers
            # seems fast enough and avoids constant allocation and
            # de-allocation which seems slower. These patches implement that
            # search.
            (0x225B99, "0c496bc4504c8bb403980600004d89f290"),
            (0x225BB8, "68"),
            (0x225BC5, "5b"),
            (
                0x225C21,
                "4183fc0675194d8b76104d85f675074c8bb3780800004d39d60f856affffff4531ffebb1",
            ),
        ),
    ),
    "game/bin/linuxsteamrt64/libscenesystem.so": patch(
        "6d395e42a86d352eed03a794bb877f3fd21e84d4be494daa309ad986208e5c48",
        "7e47f2c0ba97e60cbedcb154df4017feaf04ad1d2528b942f46da15b7303bbc6",
        (
            # Only advance the instance count after a successful allocation.
            (0x163C41, "0f86ad0000006690"),
            (0x163CF4, "89b30c5f0000e9a1ffffff"),
            # Skip this batch of instances if allocation fails.
            (0x25BA30, "e929b8ffff90"),
            (0x25725E, "89853cb7ffff4883bd58b7ffff000f844d490000e9bf470000"),
        ),
    ),
    "game/bin/linuxsteamrt64/libnetworksystem.so": patch(
        "adab8ab904bb6e751ab98943a707e9df0e3f5cd2a084ff223bb563e8ef91f08d",
        "cc4d9e260f9ead33b764c1e56f4f0ed03a87ef1a98b54149008f1b14d5b591e9",
        (
            # Use the 8-byte snappy fast path for offsets below 16.
            (0x1FAEB5, "0f"),
            (0x1FAEBA, "86"),
        ),
    ),
}


def sha256(data):
    return hashlib.sha256(data).digest()


def notify(message):
    print(f"{NAME}: {message}", file=sys.stderr)
    try:
        subprocess.run(
            ["notify-send", NAME, message],
            check=False,
            timeout=3,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
    except (OSError, subprocess.TimeoutExpired):
        pass


def replace(path, data):
    mode = stat.S_IMODE(path.stat().st_mode)
    with tempfile.NamedTemporaryFile(
        dir=path.parent, prefix=path.name + ".", delete=False
    ) as f:
        temporary = Path(f.name)
        try:
            f.write(data)
            f.flush()
            os.fchmod(f.fileno(), mode)
            os.fsync(f.fileno())
            os.replace(temporary, path)
        finally:
            temporary.unlink(missing_ok=True)


def main(command):
    diagnostics = []
    if not command:
        print(f"Usage: {NAME} command [args...]", file=sys.stderr)
        return 2
    try:
        root = Path.cwd()
        pending = []
        for relative, patch in PATCHES.items():
            path = (root / relative).resolve()
            try:
                original = path.read_bytes()
            except FileNotFoundError as e:
                diagnostics.append(f"Skipped `{relative}`: {e}")
                continue
            try:
                updated = patch.apply_to(original)
            except ValueError as e:
                raise ValueError(f"Patching `{relative}` failed: {e}") from e
            except LookupError as e:
                diagnostics.append(f"Skipped `{relative}`: {e}")
                continue
            if updated != original:
                pending.append((path, updated))
        for path, updated in pending:
            replace(path, updated)
            diagnostics.append(f"Patched {path}")
        if diagnostics:
            notify("\n".join(diagnostics))
    except (OSError, ValueError) as e:
        notify("\n".join([f"Encountered an error: {e}", *diagnostics]))
        return 1
    os.execvp(command[0], command)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))