"""Make `tile_kernels.torch.cast` free of CPU-GPU syncs.

The engine guards the forward pass with `forbid_cuda_sync()`, which turns torch's sync debug mode
to `"error"`; a single blocking copy raises `RuntimeError: called a synchronizing CUDA operation`.
The compiled kernels honour it. The reference does not, in four places:

    torch.where(mask, abs_f, torch.tensor(-1.0, device=device, dtype=...))
    torch.where(cond,     torch.tensor(0.0,  device=device), ...)

Building a tensor *on the device from a Python scalar* is a host-to-device copy, so it syncs. It is
also unnecessary: `torch.where` takes a plain scalar as its value operand and passes it to the
kernel as an argument, with no copy and no sync.

This is a defect in the reference independent of my bridge -- any caller that holds the device
clean across a cast hits it -- so the right fix is here rather than a guard lifted in the compat
layer. Re-running is a no-op.

    python patch_torch_ref_sync.py <tile_kernels-dir>
    python patch_torch_ref_sync.py <tile_kernels-dir> --revert
"""

from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

# Each replacement is (old, new, why). The scalar form is not a micro-optimisation: it is the
# difference between raising and not raising under the engine's sync guard.
REPLACEMENTS = [
    (
        "torch.where(reshaped_mask, abs_f, torch.tensor(-1.0, device=device, dtype=abs_f.dtype))",
        "torch.where(reshaped_mask, abs_f, -1.0)",
        "pad sentinel as a scalar operand",
    ),
    (
        "torch.where(dequant_sf_rounded == 0, torch.tensor(0.0, device=device), 1.0 / dequant_sf_rounded)",
        "torch.where(dequant_sf_rounded == 0, 0.0, 1.0 / dequant_sf_rounded)",
        "zero guard as a scalar operand",
    ),
    (
        "torch.where(ds_int_rounded == 0, torch.tensor(0.0, device=device), 1.0 / dequant_sf_e4m3)",
        "torch.where(ds_int_rounded == 0, 0.0, 1.0 / dequant_sf_e4m3)",
        "zero guard as a scalar operand (e4m3 sf)",
    ),
    (
        "torch.where(ds_int_rounded == 0, torch.tensor(0.0, device=device), max_quant_val_expanded / max_val)",
        "torch.where(ds_int_rounded == 0, 0.0, max_quant_val_expanded / max_val)",
        "zero guard as a scalar operand (sqr variant)",
    ),
]


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("target", type=Path, help="the installed tile_kernels package directory")
    parser.add_argument("--revert", action="store_true")
    args = parser.parse_args()

    path = args.target / "torch" / "cast.py"
    if not path.is_file():
        print(f"not found: {path}", file=sys.stderr)
        return 1
    backup = path.with_suffix(".py.orig")

    if args.revert:
        if not backup.is_file():
            print(f"no backup for {path}", file=sys.stderr)
            return 1
        shutil.move(str(backup), str(path))
        print(f"reverted {path.name}")
        return 0

    text = path.read_text()
    applied = 0
    for old, new, why in REPLACEMENTS:
        if old in text:
            text = text.replace(old, new)
            print(f"applied: {why}")
            applied += 1

    if applied == 0:
        # Either already patched, or the file changed upstream. Distinguish the two, because the
        # first is fine and the second needs a human.
        if "torch.tensor(-1.0, device=device" not in text and "torch.tensor(0.0, device=device)" not in text:
            print("already patched; nothing to do")
            return 0
        print(
            "the device-scalar pattern is still present but the exact lines changed; "
            "update REPLACEMENTS against the installed version",
            file=sys.stderr,
        )
        return 1

    if not backup.is_file():
        shutil.copy2(path, backup)
    path.write_text(text)
    print(f"\npatched {path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())