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

"""A minimally intrusive Frame SteamOS patcher"""

import argparse
import hashlib
import json
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

MARKER = "x-valve-please-fix"
RUNTIME = Path(os.environ.get("RUNTIME_DIRECTORY", "/run/patcher"))


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


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

    def apply_to(self, data: bytes) -> bytearray:
        result = bytearray(data)
        for edit in self.edits:
            end = edit.offset + len(edit.replacement)
            if edit.offset < 0 or end > len(result):
                raise ValueError(f"Patch outside file at {edit.offset:#x}")
            result[edit.offset : end] = edit.replacement
        return result


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


PATCHES = {
    "/opt/steamvr/tools/vrlink/bin/linuxarm64/vrlink": patch(
        original_hash="31addef0c9553d3770b5570487699e96d70271d4073cd6c03b080fdee6ae2153",
        edits=(
            # Clear the lists of ready and submitted frame IDs after a reset.
            # If not done, this can cause empty frames to be submitted to the
            # decoder, or skip submitting a received frame that hasn't been
            # submitted since the reset.
            (0x53EB4, "75d40214"),
            (
                0x109088,
                "6072049161720c9102008092020881a81f0001ebc1ffff54616202916082249103048052020881a863040071c1ffff54802bfd17",
            ),
            # Let the decoder worker continue running when either DQBUF returns
            # EAGAIN. This can happen after a reset, and will trigger more
            # resets, and can only recover after the decoder stalls for at
            # least an absurd 0.8s.
            (0x32D50, "db580314"),
            (0x1090BC, "090040b93f2d0071e0d29454130040b922a7fc17"),
            (0x32D90, "d0580314"),
            (0x1090D0, "090040b93f2d007140d89454130040b92da7fc17"),
            # Fix a possible bug where a missing frame wouldn't trigger a
            # request when a newer frame arrived and the gap straddled the 16
            # bit wrapping boundary.
            (0x53F84, "a102024b213c00133f080071ab0100547316931a06000014"),
            (0x53FB0, "613e0053"),
            # Ignore packets which are for an older frame when re-assembling.
            # There are 16 slots and this could otherwise cause an old packet
            # to override a newer one.
            (0x53F34, "6cd40214"),
            (0x1090E4, "c00a40b96000f837a002004b60007837c0324039902bfd17b22bfd17"),
            # Avoid two corner cases:
            # - After reset, wait for a complete recovery frame instead of
            #   simply waiting for _a_ frame to avoid trying to start decoding
            #   from a frame which depends on prior ones.
            # - Ignore frames which complete reassembly but which are older
            #   than the last submitted frame.
            (0x541E0, "c8d30214"),
            (
                0x109100,
                "092061f92a1941b9ca00f8370b0840b96b010a4b0b017837fd7bbea9322cfd170a3040398a0000340a1040914a754d394aff0f37c0035fd6",
            ),
            # To compensate for the first of the above two fixes, any frames
            # which completed out of order before a recovery frame must be
            # revisited after that frame.
            (0x53A90, "41b55a54"),
            (
                0x109138,
                "800e4092000f008b004c40f9010840b93f00146be150a554242cfd97852afd17",
            ),
            # If decoder submission fails, request reset without advancing the
            # last-submitted frame number as though submission succeeded.
            (0x53B08, "94d50214"),
            (0x109158, "4050a534001b46f96b2afd17"),
            # With FEC, audio gets sent as two blocks (A, B) followed by their
            # parity (A^B) block. If the next block arrives (C) while the
            # headset has received B but not A and is still waiting for A^B
            # then B gets discarded without playing it. This loses two blocks
            # worth of audio (10.67ms) rather than just skipping one missed
            # block (5.33ms) and flushing the second one (B). This fixes that
            # issue by just flushing B in that case.
            (0x3ADA8, "ef380314"),
            (
                0x109164,
                "2000004b6000f836e2ff9f520000020b1f080071ece19854cbe49854620080522308c21a6384021b43e1983421c7fc17",
            ),
            # Patch incorrect sequence gap calculation when sequence numbers
            # wrap modulo 65535 (recurring theme...)
            (0x3AFB0, "79380314"),
            (0x109194, "6000f836e1ff9f520000010b0dc7fc17"),
            # Don't request retransmission of frames that the headset has
            # already moved ahead of (i.e. the last-submitted frame or older).
            # Assume that a frame with a size of 4 MiB with missing metadata is
            # not actually 4 MiB to avoid requesting way too many potentially
            # missing fragments automatically. This can cause too many requests
            # which triggers an unnecessary stream reset. Requests for the
            # frame's metadata fragment continue which will contain and
            # populate the size field allowing the retransmit request code to
            # work normally from that point.
            (0x52BF8, "6bd90214"),
            (
                0x1091A4,
                "832261f9641841b99f04003180000054e400044bc4e3a434e40078378032403960e2a435951240b9bf025071f503951a8d26fd171627fd17",
            ),
            # Don't trigger a reset if more than 512 missing fragments need to
            # be requested as a result of a time based repair. Instead stop
            # early once 20 entries are accumulated, reserving an extra 21st
            # entry so that a request for a fresh metadata fragment can always
            # be sent. This avoids an unnecessary reset and can speed up
            # recovery.
            (0x52E8C, "d4d80214"),
            (
                0x1091DC,
                "130080523f04003161e5a45484364079a0fe0f11007c0a539f00006bc2e4a45402008012e2c30079330080522227fd17",
            ),
            (0x52EC8, "d1d80214"),
            (0x10920C, "8a3240394a01130a4a4d0011852640b92c27fd17"),
            (0x52F38, "7f020a6b6c040054"),
            (0x52F98, "7f020a6b6c010054"),
            (0x52FC0, "7f020a6bedeaff5453050011647a1f53"),
        ),
    ),
}


def start():
    for name, patch in PATCHES.items():
        target = Path(name).resolve(strict=True)
        data = target.read_bytes()
        checksum = hashlib.sha256(data).digest()
        if checksum != patch.original_hash:
            print(
                f"Skipping {target}: unrecognised checksum {checksum.hex()}",
                file=sys.stderr,
            )
            continue
        patched = patch.apply_to(data)
        source = RUNTIME / hashlib.sha256(os.fsencode(target)).hexdigest()
        source.parent.mkdir(parents=True, exist_ok=True)
        source.write_bytes(patched)
        attributes = target.stat()
        os.chown(source, attributes.st_uid, attributes.st_gid)
        source.chmod(attributes.st_mode)
        subprocess.run(
            ["mount", "-o", f"bind,ro,{MARKER}", str(source), str(target)],
            check=True,
        )
        print(f"Patched {target}", file=sys.stderr)


def stop():
    result = subprocess.run(
        [
            "findmnt",
            "--mtab",
            "--options",
            MARKER,
            "--json",
            "--output",
            "TARGET",
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    for target in reversed(
        [
            Path(row["target"])
            for row in json.loads(result.stdout).get("filesystems", [])
        ]
    ):
        subprocess.run(["umount", str(target)], check=True)
        print(f"Unmounted {target}", file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=("start", "stop"))
    args = parser.parse_args()
    if args.action == "start":
        start()
    else:
        stop()


if __name__ == "__main__":
    main()