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

"""A steamvr launch command wrapper that patches various parts of steamvr
to fix bugs.

For SteamVR Beta 2.19.1, build 26389675.

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 a 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 = {
    "bin/linux64/vrcompositor": patch(
        "4c380adc27832db9afaf83e446080354338bcb4c96f0f0cc561c90cd4d79168f",
        "a6d073c17489919e838c3d2ac4941a923bc18172a372f3c32bae3433e95f9dbc",
        (
            # Wait for GPU to finish using its resources (see #952)
            (0x404941, "84c00f842f570000eb06"),
            (0x404951, "488dbbe00d0000e90a570000"),
            (0x40A05B, "e830db0600"),
            (0x40A060, "e9dca8ffff9090"),
            # Stop reusing GPU work identifier across frames
            # (allows above patch to work reliably)
            (0x41D224, "31c09090909090"),
            (0x41DB39, "31c09090909090"),
            (0x41DB74, "909090909090"),
            # Fix shader descriptor texture array counts (should be 3, not 1)
            # for bindings 9, 17, and 20. Affected shaders:
            # - distort_none_ps
            # - distort_none_ps_layered
            # - distort_panorama_ps
            # - distort_panorama_solid_color_ps
            # - distort_ps_achromatic_nd
            (0x404961, "498b5500b80300000089424ceb22"),
            (0x404991, "89828c0000008982a4000000eb12"),
            (0x4049B1, "4c8bbd18feffffe902200100"),
            (0x4169B8, "e9a4dffeff9090"),
        ),
    ),
    "bin/linux64/vrserver": patch(
        "a7c9246f8fe36749caa0bfe3984efbb5fa7dbbb28e5263a5c1744f5ba6b0c10b",
        "b14314f1fbfe8acb164583853e32c9a6870317bb9cbb278d9cea21862aaebd8f",
        (
            # Remove logger calls in forked processes (can cause a deadlock and
            # make the process linger preventing SteamVR from starting again)
            (0x2C1386, "9090909090"),
            (0x2C1577, "9090909090"),
            (0x2C160C, "9090909090"),
            (0x2C1726, "9090909090"),
            (0x2C17FF, "9090909090"),
            (0x2C1699, "9090909090"),
        ),
    ),
    "drivers/vrlink/bin/linux64/driver_vrlink.so": patch(
        "a33d1da2d1c9851a61e19b87ffc89bbbb6e8b1e503e189e3e19bccd2e6cd3b7f",
        "c7d8adbbcd44d577f7891c15e53d6000eaf0b257fde642b09ff286bd00be9468",
        (
            # Keep increasing (mod(2^16)) the ACK cutoff instead of clearing it
            # (#965)
            (0xD91DD, "89836c0d000090909090"),
            # Wait for encoder to finish before reconfiguring it
            # (ulVideoTimeline == ulTimelineSequence) (#965)
            (0x13C2D2, "83bb4c090000007458e9702c0000"),
            (0x13C333, "488b8328090000e9ec210000"),
            (0x13E524, "e9a9ddffff9090"),
            # Clean up failed encoded frames so the above fix doesn't wait
            # forever when these trigger a reset
            (0x13D735, "e927000000"),
            (0x13D761, "488b9d70ffffff4c01f3e9d6feffff"),
            # Reset tracking and encoding when starting a new streaming session
            (0x1129F1, "57488b7f184885ff740feb04"),
            (0x112A01, "488b076a015eff50285fe91013feff"),
            (0x11A251, "4c89a30804000031c0894308c3"),
            (0x11A297, "e8b5ffffff6690"),
            (0x11A2AA, "e84287ffff"),
            # Ignore stale tracking IDs and repeated older attempts
            (0x11B301, "89d02b47084189c6c1e004"),
            (0x11B30C, "0f8853050000"),
            (0x11B614, "0f841719fdff"),
            (0xECF31, "410fb644240988430ee926e90200"),
        ),
    ),
}


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:]))