from pathlib import Path
import os
import subprocess

# The upstream checkout lives on persistent storage and therefore retains edits
# across container restarts. Restore only the files managed by this patcher from
# its pinned Git commit before applying changes, making this script idempotent.
for relative_path in (
    "util/dataset_mixamo_additional.py",
    "model.py",
    "models_ae.py",
    "app.py",
):
    pristine = subprocess.check_output(
        ["git", "-C", "app", "show", f"HEAD:{relative_path}"]
    )
    Path("app", relative_path).write_bytes(pristine)

# ── 1. dataset_mixamo_additional.py — optional extended skeleton guard ─────────
path = Path("app/util/dataset_mixamo_additional.py")
if path.exists():
    source = path.read_text()
    old = """TEMPLATE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/Mixamo/bones_vroid.fbx"))
KINEMATIC_TREE = get_kinematic_tree(TEMPLATE_PATH, BONES_IDX_DICT)
assert len(KINEMATIC_TREE) == len(MIXAMO_JOINTS)
"""
    new = """TEMPLATE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/Mixamo/bones_vroid.fbx"))
if os.path.isfile(TEMPLATE_PATH):
    KINEMATIC_TREE = get_kinematic_tree(TEMPLATE_PATH, BONES_IDX_DICT)
    assert len(KINEMATIC_TREE) == len(MIXAMO_JOINTS)
else:
    # The optional extended skeleton lives in the restricted Mixamo dataset.
    # Standard humanoid inference does not use it (joints_additional=False).
    KINEMATIC_TREE = dataset_mixamo.KINEMATIC_TREE
"""
    if old in source:
        path.write_text(source.replace(old, new))
        print("[dataset_mixamo_additional.py] extended-skeleton guard patched")
    else:
        print("[dataset_mixamo_additional.py] already patched or block not found")

# ── fps fallback: replaces torch_cluster.fps with pytorch3d equivalent ─────────
FPS_FALLBACK = """try:
    from torch_cluster import fps
except ImportError:
    import numpy as np
    from pytorch3d.ops import sample_farthest_points

    def fps(x, batch, ratio):
        import torch
        sampled = []
        for batch_id in torch.unique(batch, sorted=True):
            global_idx = torch.nonzero(batch == batch_id, as_tuple=False).flatten()
            points = x[global_idx].unsqueeze(0)
            count = max(1, int(np.ceil(points.shape[1] * ratio)))
            _, local_idx = sample_farthest_points(points, K=count, random_start_point=True)
            sampled.append(global_idx[local_idx[0]])
        return torch.cat(sampled)
"""
CLUSTER_IMPORT = "from torch_cluster import fps"

# ── 2. model.py ───────────────────────────────────────────────────────────────
model_path = Path("app/model.py")
if model_path.exists():
    src = model_path.read_text()
    if CLUSTER_IMPORT in src:
        model_path.write_text(src.replace(CLUSTER_IMPORT, FPS_FALLBACK, 1))
        print("[model.py] torch_cluster fps patched")
    else:
        print("[model.py] already patched")
else:
    print("[model.py] not found")

# ── 3. models_ae.py — THE CRITICAL MISSING PATCH ─────────────────────────────
# The original patch_runtime.py only patched model.py but NOT models_ae.py,
# which also imports torch_cluster. This caused the ModuleNotFoundError.
models_ae_path = Path("app/models_ae.py")
if models_ae_path.exists():
    src = models_ae_path.read_text()
    if CLUSTER_IMPORT in src:
        models_ae_path.write_text(src.replace(CLUSTER_IMPORT, FPS_FALLBACK, 1))
        print("[models_ae.py] torch_cluster fps patched")
    else:
        print("[models_ae.py] already patched")
else:
    print("[models_ae.py] not found")

# ── 4. app/app.py — @spaces.GPU duration ─────────────────────────────────────
app_path = Path("app/app.py")
if app_path.exists():
    src = app_path.read_text()
    bare_gpu_decorators = src.count("@spaces.GPU\n")
    if bare_gpu_decorators != 2:
        raise RuntimeError(f"Expected 2 bare GPU decorators, found {bare_gpu_decorators}")
    app_path.write_text(src.replace("@spaces.GPU\n", "@spaces.GPU(duration=300)\n"))
    print(f"[app/app.py] {bare_gpu_decorators} @spaces.GPU -> duration=300 patched")
else:
    print("[app/app.py] not found")

print("patch_runtime.py done.")