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",
    "app_blender.py",
    "util/blender_utils.py",
):
    pristine = subprocess.check_output(
        ["git", "-c", "safe.directory=/data/app", "-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.")


# STATICAI Standard animation library integration
def _replace_once(path, old, new, label):
    source = path.read_text()
    count = source.count(old)
    if count != 1:
        raise RuntimeError(f"{label}: expected exactly one patch target, found {count}")
    path.write_text(source.replace(old, new, 1))


blender_utils_path = Path("app/util/blender_utils.py")
_replace_once(
    blender_utils_path,
    "def load_mixamo_anim(char_file: str, anim_file: str, do_retarget=False, inplace=False, to_tris=False):",
    "def load_mixamo_anim(char_file: str, anim_file: str, do_retarget=False, inplace=False, to_tris=False, action_name=None):",
    "blender_utils signature",
)
_replace_once(
    blender_utils_path,
    '    assert anim_armature.animation_data is not None and len(bpy.data.actions) > 0, f"Animation not found in {anim_file}"\n\n    set_action(char_armature, anim_armature.animation_data.action)',
    '''    assert anim_armature.animation_data is not None and len(bpy.data.actions) > 0, f"Animation not found in {anim_file}"

    if action_name:
        candidates = [
            action
            for action in bpy.data.actions
            if action.name == action_name
            or action.name.endswith("|" + action_name)
            or action.name.startswith(action_name + ".")
        ]
        if not candidates:
            available = ", ".join(sorted(action.name for action in bpy.data.actions))
            raise ValueError(f"Animation action {action_name!r} not found. Available: {available}")
        anim_armature.animation_data.action = candidates[0]
        print(f"Selected animation action: {candidates[0].name}")

    set_action(char_armature, anim_armature.animation_data.action)''',
    "blender_utils action selection",
)

app_blender_path = Path("app/app_blender.py")
_replace_once(
    app_blender_path,
    '''            blender_utils.load_mixamo_anim(
                [armature_obj, mesh_obj], args.animation_path, do_retarget=args.retarget, inplace=args.inplace
            )''',
    '''            blender_utils.load_mixamo_anim(
                [armature_obj, mesh_obj],
                args.animation_path,
                do_retarget=args.retarget,
                inplace=args.inplace,
                action_name=args.animation_action,
            )''',
    "app_blender load_mixamo_anim",
)
_replace_once(
    app_blender_path,
    '    parser.add_argument("--animation_path", type=str, default=None)\n',
    '    parser.add_argument("--animation_path", type=str, default=None)\n    parser.add_argument("--animation_action", type=str, default=None)\n',
    "app_blender parser action",
)

app_path = Path("app/app.py")
_replace_once(app_path, "import gc\nimport os\n", "import gc\nimport json\nimport os\n", "app json import")
_replace_once(
    app_path,
    '''    animation_file: str,
    retarget: bool,''',
    '''    animation_file: str,
    animation_action: str,
    retarget: bool,''',
    "vis_blender action argument",
)
_replace_once(
    app_path,
    '''                animation_path=animation_file,
                retarget=retarget,''',
    '''                animation_path=animation_file,
                animation_action=animation_action,
                retarget=retarget,''',
    "Namespace action argument",
)
_replace_once(
    app_path,
    '''            if animation_file is not None:
                cmd += f" --animation_path '{os.path.abspath(animation_file)}'"
                if retarget:
                    cmd += " --retarget"
                if inplace:
                    cmd += " --inplace"''',
    '''            if animation_file is not None:
                cmd += f" --animation_path '{os.path.abspath(animation_file)}'"
                if animation_action:
                    safe_action = animation_action.replace("'", "")
                    cmd += f" --animation_action '{safe_action}'"
                if retarget:
                    cmd += " --retarget"
                if inplace:
                    cmd += " --inplace"''',
    "subprocess action argument",
)
_replace_once(
    app_path,
    '''    animation_file: str = None,
    retarget=True,''',
    '''    animation_file: str = None,
    animation_action: str = None,
    retarget=True,''',
    "pipeline action argument",
)
_replace_once(
    app_path,
    '''    yield vis_blender(
        reset_to_rest, no_fingers, rest_pose_type, ignore_pose_parts, animation_file, retarget, inplace, db
    )''',
    '''    yield vis_blender(
        reset_to_rest,
        no_fingers,
        rest_pose_type,
        ignore_pose_parts,
        animation_file,
        animation_action,
        retarget,
        inplace,
        db,
    )''',
    "pipeline vis_blender action",
)
_replace_once(app_path, '    title = "Make-It-Animatable"\n', '    title = "Motion Rig"\n', "Motion Rig title")
_replace_once(
    app_path,
    '    animation_dir = "data/Mixamo/animation"\n',
    '''    animation_dir = "data/Mixamo/animation"
    standard_library_root = "/data/animations/Standard"
    standard_manifest_path = os.path.join(standard_library_root, "ANIMATION_MANIFEST.json")
    standard_animation_lookup = {}
    standard_animation_choices = []
    if os.path.isfile(standard_manifest_path):
        with open(standard_manifest_path, "r", encoding="utf-8") as fh:
            manifest = json.load(fh)
        for entry_index, entry in enumerate(manifest.get("entries", [])):
            motion_file = os.path.join(standard_library_root, entry["file"])
            variant = "Root Motion" if entry.get("root_motion") else "In Place"
            for action in entry.get("actions", []):
                if not action or action == "A_TPose":
                    continue
                key = f"{entry_index}:{action}"
                label = f"Standard {entry.get('set')} | {action} | {variant}"
                standard_animation_lookup[key] = {"path": motion_file, "action": action}
                standard_animation_choices.append((label, key))
''',
    "Standard library manifest load",
)
_replace_once(
    app_path,
    '                                    sorted(glob("*.fbx", root_dir=animation_dir)),\n',
    '                                    standard_animation_choices,\n',
    "Standard animation choices",
)
_replace_once(
    app_path,
    '                                    label="Select Animation File",\n',
    '                                    label="Standard Animation Library",\n',
    "Standard animation label",
)
_replace_once(
    app_path,
    '''                                input_animation_file = gr.File(
                                    label="Animation File",
                                    file_types=[".fbx"],
                                    value=lambda: "./data/Standard Run.fbx",
                                    interactive=True,
                                )''',
    '''                                input_animation_file = gr.File(
                                    label="Animation File",
                                    file_types=[".fbx"],
                                    value=lambda: "./data/Standard Run.fbx",
                                    interactive=True,
                                )
                                animation_action_state = gr.State(value=None)''',
    "animation action state",
)
_replace_once(
    app_path,
    '''                input_animation_file,
                input_retarget,''',
    '''                input_animation_file,
                animation_action_state,
                input_retarget,''',
    "inputs tuple action state",
)
_replace_once(
    app_path,
    '''                        animation_file=inputs[input_animation_file],
                        retarget=inputs[input_retarget],''',
    '''                        animation_file=inputs[input_animation_file],
                        animation_action=inputs[animation_action_state],
                        retarget=inputs[input_retarget],''',
    "pipeline UI action",
)
_replace_once(
    app_path,
    '''                    input_animation_file,
                    input_retarget,''',
    '''                    input_animation_file,
                    animation_action_state,
                    input_retarget,''',
    "animate inputs action state",
)
_replace_once(
    app_path,
    '''            def select2file(selected: str):
                return None if selected is None else os.path.join(animation_dir, selected)

            select_animation_file.input(
                lambda x: (select2file(x), gr.Checkbox(value=True)),
                inputs=select_animation_file,
                outputs=[input_animation_file, input_reset_to_rest],
                preprocess=False,
            )
            input_animation_file.upload(
                lambda: (gr.Dropdown(value=None), gr.Checkbox(value=True)),
                outputs=[select_animation_file, input_reset_to_rest],
            )
            input_animation_file.clear(
                lambda: gr.Dropdown(value=None), outputs=select_animation_file, show_progress="hidden"
            )''',
    '''            def select2file(selected: str):
                if selected is None:
                    return None, None
                item = standard_animation_lookup.get(selected)
                if item is None:
                    raise gr.Error(f"Unknown Standard animation selection: {selected}")
                return item["path"], item["action"]

            def select_library_animation(selected: str):
                path, action = select2file(selected)
                return path, action, gr.Checkbox(value=True)

            select_animation_file.input(
                select_library_animation,
                inputs=select_animation_file,
                outputs=[input_animation_file, animation_action_state, input_reset_to_rest],
                preprocess=False,
            )
            input_animation_file.upload(
                lambda: (gr.Dropdown(value=None), None, gr.Checkbox(value=True)),
                outputs=[select_animation_file, animation_action_state, input_reset_to_rest],
            )
            input_animation_file.clear(
                lambda: (gr.Dropdown(value=None), None),
                outputs=[select_animation_file, animation_action_state],
                show_progress="hidden",
            )''',
    "Standard selection events",
)

print("STATICAI Standard animation library integration patched.")