#!/usr/bin/env python3
"""Patch an installed lerobot==0.5.1 to be Python 3.11-compatible.

lerobot 0.5.1 uses PEP 695 syntax (generic class/function parameter lists
and the `type` alias statement, both 3.12-only) at 4 sites. Python 3.11's
parser rejects all of them. This script rewrites them to pre-PEP-695
forms (typing.TypeVar + Generic, plain assignments) so the package can
be imported under Python 3.11.

Idempotent: safe to run multiple times. Verifies all touched files parse
under Python 3.11 grammar at the end (fails loud if anything regressed).

Use in a Dockerfile:
    RUN pip install --ignore-requires-python "lerobot==0.5.1" && \\
        python /tmp/patch_lerobot_for_py311.py

Why --ignore-requires-python: lerobot pins requires-python>=3.12 in its
metadata. After this patch the code itself runs on 3.11; the pin is no
longer accurate for the patched copy.
"""
from __future__ import annotations

import ast
import importlib.util
import sys
from pathlib import Path


def _locate_lerobot() -> Path:
    spec = importlib.util.find_spec("lerobot")
    if spec is None or spec.origin is None:
        sys.exit("error: lerobot is not installed in this environment")
    return Path(spec.origin).parent


def _replace_in_file(path: Path, replacements: list[tuple[str, str]]) -> bool:
    """Apply each (old, new) replacement to `path`. Returns True if any
    change was made. Idempotent: if `new` is already present, treat as
    already-patched even when `old` is also (substring-)present, since
    several of our replacements append to a block where `old` is a
    prefix of `new`."""
    src = path.read_text()
    original = src
    for old, new in replacements:
        if new in src:
            continue  # already patched
        if old in src:
            src = src.replace(old, new, 1)  # only the first match
        else:
            sys.exit(
                f"error: {path} does not contain expected old or new text:\n"
                f"  old: {old!r}\n  new: {new!r}"
            )
    if src != original:
        path.write_text(src)
        return True
    return False


def patch(root: Path) -> None:
    print(f"patching lerobot at {root}")

    # 1) motors/motors_bus.py — drop `type` keyword on two alias statements.
    changed = _replace_in_file(
        root / "motors" / "motors_bus.py",
        [
            ("type NameOrID = str | int\ntype Value = int | float",
             "NameOrID = str | int\nValue = int | float"),
        ],
    )
    print(f"  motors_bus.py            : {'patched' if changed else 'already patched'}")

    # 2) datasets/streaming_dataset.py — Backtrackable[T] -> Generic[T].
    changed = _replace_in_file(
        root / "datasets" / "streaming_dataset.py",
        [
            ("from collections import deque\n"
             "from collections.abc import Callable, Generator, Iterable, Iterator\n"
             "from pathlib import Path",
             "from collections import deque\n"
             "from collections.abc import Callable, Generator, Iterable, Iterator\n"
             "from pathlib import Path\n"
             "from typing import Generic, TypeVar\n\n"
             "T = TypeVar(\"T\")"),
            ("class Backtrackable[T]:", "class Backtrackable(Generic[T]):"),
        ],
    )
    print(f"  streaming_dataset.py     : {'patched' if changed else 'already patched'}")

    # 3) utils/io_utils.py — function generic + JsonLike-bound T.
    changed = _replace_in_file(
        root / "utils" / "io_utils.py",
        [
            ("import json\nimport warnings\nfrom pathlib import Path\n\n"
             "import imageio\n\n"
             "JsonLike = ",
             "import json\nimport warnings\nfrom pathlib import Path\n"
             "from typing import TypeVar\n\n"
             "import imageio\n\n"
             "JsonLike = "),
            # Inject the bound TypeVar right after the JsonLike alias.
            ("JsonLike = str | int | float | bool | None | "
             "list[\"JsonLike\"] | dict[str, \"JsonLike\"] | tuple[\"JsonLike\", ...]\n\n\n"
             "def write_video",
             "JsonLike = str | int | float | bool | None | "
             "list[\"JsonLike\"] | dict[str, \"JsonLike\"] | tuple[\"JsonLike\", ...]\n"
             "T = TypeVar(\"T\", bound=JsonLike)\n\n\n"
             "def write_video"),
            ("def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:",
             "def deserialize_json_into_object(fpath: Path, obj: T) -> T:"),
        ],
    )
    print(f"  io_utils.py              : {'patched' if changed else 'already patched'}")

    # 4) processor/pipeline.py — DataProcessorPipeline[TInput, TOutput] -> Generic.
    changed = _replace_in_file(
        root / "processor" / "pipeline.py",
        [
            ("from typing import Any, TypedDict, TypeVar, cast",
             "from typing import Any, Generic, TypedDict, TypeVar, cast"),
            ("class DataProcessorPipeline[TInput, TOutput](HubMixin):",
             "class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):"),
        ],
    )
    print(f"  processor/pipeline.py    : {'patched' if changed else 'already patched'}")


def verify(root: Path) -> None:
    print("verifying all lerobot .py files parse under Python 3.11 grammar...")
    fails: list[tuple[Path, int, str]] = []
    for p in root.rglob("*.py"):
        try:
            ast.parse(p.read_text(), filename=str(p), feature_version=(3, 11))
        except SyntaxError as e:
            fails.append((p, e.lineno or 0, e.msg or ""))
    if fails:
        print(f"  ✗ {len(fails)} file(s) still fail 3.11 parse:")
        for p, ln, msg in fails:
            print(f"    {p.relative_to(root)}:{ln} {msg}")
        sys.exit(1)
    print("  ✓ all files compatible with Python 3.11")


def main() -> None:
    root = _locate_lerobot()
    patch(root)
    verify(root)


if __name__ == "__main__":
    main()