#!/usr/bin/env python3
"""Point the unpacked Quest Rec Room 20230907 client at a Radicade server.

Accepts either the raw apktool-decoded APK directory (containing
assets/bin/Data/Managed/Metadata/global-metadata.dat) or a wrapper containing
apk/ and obb/ subdirectories like the provided download. It changes only exact
IL2CPP literal-table rows, not broad byte substrings.

Dry-run by default. Use --apply to back up the metadata and write the patch.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import struct
import sys
import tempfile
from pathlib import Path
from urllib.parse import urlsplit

# The Quest release's header fields were mapped structurally from its unique
# metadata section chain. These contain the v27 string-literal table/data pairs.
FIELDS = {
    "literal_table_offset": 0x44,
    "literal_table_size": 0x6C,
    "literal_data_offset": 0x28,
    "literal_data_size": 0xEC,
}
GAME_VERSION = "20230907"
KNOWN_LITERALS = {
    2011: ".rec.net",
    2516: "20230907",
    41319: "http://cdn.rec.net/avatar/",
    41405: "https://auth.rec.net",
    41411: "https://ns.rec.net",
    41418: "https://rec.net",
    41419: "https://rec.net/shop",
    41420: "https://rec.net/shop/subscriptions",
    41421: "https://rec.net/{0}",
}


def u32(data: bytes | bytearray, offset: int) -> int:
    return struct.unpack_from("<I", data, offset)[0]


def validate_server(value: str) -> str:
    parsed = urlsplit(value)
    if parsed.scheme != "https" or not parsed.netloc or parsed.username or parsed.password:
        raise ValueError("server must be an https:// URL without user info")
    if parsed.query or parsed.fragment:
        raise ValueError("server URL cannot contain a query or fragment")
    return value.rstrip("/")


def load_literals(blob: bytes) -> tuple[int, int, int, int, list[tuple[int, int]]]:
    if len(blob) < 256:
        raise ValueError("metadata is shorter than its 256-byte header")
    table_offset = u32(blob, FIELDS["literal_table_offset"])
    table_size = u32(blob, FIELDS["literal_table_size"])
    data_offset = u32(blob, FIELDS["literal_data_offset"])
    data_size = u32(blob, FIELDS["literal_data_size"])
    if table_size % 8 or table_offset < 256 or table_offset + table_size > len(blob):
        raise ValueError("Quest metadata literal-table header fields are invalid")
    if data_offset < 256 or data_offset + data_size > len(blob):
        raise ValueError("Quest metadata literal-data header fields are invalid")

    literals = []
    for index in range(table_size // 8):
        length, string_index = struct.unpack_from("<Ii", blob, table_offset + index * 8)
        if length == 0:
            if string_index != 0:
                raise ValueError(f"empty literal #{index} has unexpected data index {string_index}")
            literals.append((length, string_index))
            continue
        if string_index < 0 or string_index + length > data_size:
            raise ValueError(f"literal #{index} points outside its data segment")
        try:
            blob[data_offset + string_index:data_offset + string_index + length].decode("utf-8")
        except UnicodeDecodeError as exc:
            raise ValueError(f"literal #{index} is not UTF-8") from exc
        literals.append((length, string_index))
    return table_offset, table_size, data_offset, data_size, literals


def build_patch(blob: bytes, server: str) -> tuple[bytearray, list[tuple[int, str, str]], int, int]:
    table_offset, table_size, data_offset, data_size, literals = load_literals(blob)
    # Refuse to patch a different client or a normalized/tooling metadata file.
    for index, expected in KNOWN_LITERALS.items():
        if index >= len(literals):
            raise ValueError(f"expected Quest 20230907 literal #{index} is missing")
        length, string_index = literals[index]
        actual = blob[data_offset + string_index:data_offset + string_index + length].decode("utf-8")
        if actual != expected:
            raise ValueError(f"not the expected Quest 20230907 metadata: literal #{index} is {actual!r}, expected {expected!r}")

    rules = {
        "https://ns.rec.net": server,
        "https://auth.rec.net": server,
        "https://rec.net/shop/subscriptions": server + "/shop/subscriptions",
        "https://rec.net/shop": server + "/shop",
        "https://rec.net/{0}": server + "/{0}",
        "https://rec.net": server,
        # The server's signed image endpoint is /images/img/{name}.
        "http://cdn.rec.net/avatar/": server + "/images/img/",
        ".rec.net": ".radicade.internal",
    }
    rows = []
    for index, (length, string_index) in enumerate(literals):
        if not length:
            continue
        old = blob[data_offset + string_index:data_offset + string_index + length].decode("utf-8")
        # Retain the server host tail used by the runtime's image signature key id.
        if index == 2011:
            replacement = ".radicade.internal"
            new = replacement if old == ".rec.net" else old
        else:
            new = old
            # Longest-first prevents the root literal from consuming the shop forms.
            for source, target in sorted(rules.items(), key=lambda item: -len(item[0])):
                if source in new:
                    new = new.replace(source, target)
        if new != old:
            rows.append((index, old, new))

    if not rows:
        raise ValueError("no target literals found; this file may already be patched or not match the supported Quest build")
    out = bytearray(blob)
    appended = bytearray()
    new_data_offsets: dict[int, int] = {}
    for index, _old, new in rows:
        encoded = new.encode("utf-8")
        new_data_offsets[index] = data_size + len(appended)
        appended.extend(encoded)
    # As on the desktop patch, keep the old literal bytes intact and append the
    # replacement strings, then update only the matching table rows and fields.
    out.extend(blob[data_offset:data_offset + data_size])
    out.extend(appended)
    new_data_size = data_size + len(appended)
    struct.pack_into("<I", out, FIELDS["literal_data_offset"], len(blob))
    struct.pack_into("<I", out, FIELDS["literal_data_size"], new_data_size)
    for index, _old, new in rows:
        row_offset = table_offset + index * 8
        struct.pack_into("<I", out, row_offset, len(new.encode("utf-8")))
        struct.pack_into("<i", out, row_offset + 4, new_data_offsets[index])

    # Read back all live string rows; unchanged rows must resolve identically.
    relocated_data_offset = len(blob)
    changed = {index: new for index, _old, new in rows}
    for index, (old_length, old_index) in enumerate(literals):
        row_offset = table_offset + index * 8
        length, string_index = struct.unpack_from("<Ii", out, row_offset)
        if length == 0:
            continue
        actual = out[relocated_data_offset + string_index:relocated_data_offset + string_index + length]
        if index in changed:
            expected = changed[index].encode("utf-8")
        else:
            expected = blob[data_offset + old_index:data_offset + old_index + old_length]
        if actual != expected:
            raise ValueError(f"read-back verification failed for literal #{index}")
    return out, rows, data_size, new_data_size


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("apktool_dir", help="apktool-decoded directory containing apk/assets/bin/Data")
    parser.add_argument("--server", default="https://ns.radicade.app")
    parser.add_argument("--apply", action="store_true", help="write files (dry-run is the default)")
    args = parser.parse_args()
    try:
        server = validate_server(args.server)
    except ValueError as exc:
        print(f"invalid server URL: {exc}", file=sys.stderr)
        return 2

    root = Path(args.apktool_dir).resolve()
    candidates = [
        root / "assets" / "bin" / "Data" / "Managed" / "Metadata" / "global-metadata.dat",
        root / "apk" / "assets" / "bin" / "Data" / "Managed" / "Metadata" / "global-metadata.dat",
    ]
    metadata = next((candidate for candidate in candidates if candidate.is_file()), candidates[-1])
    if not metadata.is_file():
        print(f"decoded Quest project missing {metadata}", file=sys.stderr)
        return 1

    try:
        source = metadata.read_bytes()
        patched, rows, old_size, new_size = build_patch(source, server)
    except (OSError, ValueError, struct.error) as exc:
        print(f"cannot patch Quest metadata: {exc}", file=sys.stderr)
        return 1

    print(f"metadata: {metadata}")
    print(f"source sha256: {hashlib.sha256(source).hexdigest()}")
    print(f"metadata size: {len(source):,} -> {len(patched):,}; literal bytes: {old_size:,} -> {new_size:,}")
    print(f"server: {server}")
    print(f"{len(rows)} live literal(s) will be redirected:")
    for index, before, after in rows:
        print(f"  #{index}: {before!r}\n      -> {after!r}")
    print("verification: all literal rows read back; the patch changes only the table rows and its two data-segment header fields")
    if not args.apply:
        print("dry-run only; pass --apply to write the metadata backup and marker")
        return 0

    backup = Path(str(metadata) + ".orig")
    if backup.exists() and backup.read_bytes() != source:
        print(f"refusing to overwrite existing backup with different bytes: {backup}", file=sys.stderr)
        return 1
    if not backup.exists():
        shutil.copy2(metadata, backup)
    # Atomic replacement prevents a partially written 30 MB metadata file.
    fd, temp_name = tempfile.mkstemp(prefix=metadata.name + ".", dir=metadata.parent)
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(patched)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temp_name, metadata)
    except Exception:
        try:
            os.unlink(temp_name)
        except OSError:
            pass
        raise

    marker = metadata.parents[2] / "radicade-quest-patch.json"
    marker.write_text(json.dumps({
        "gameVersion": GAME_VERSION,
        "serverUrl": server,
        "sourceSha256": hashlib.sha256(source).hexdigest(),
        "patchedSha256": hashlib.sha256(patched).hexdigest(),
        "changedLiterals": len(rows),
    }, indent=2) + "\n", encoding="utf-8")
    print(f"backup: {backup}\nmarker: {marker}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())