"""
Fix material bindings on repackaged test tube USD.

The referenced meshes have bindings pointing to /Test_Tube_AA_01/Looks/...
which is outside the reference scope. This script overrides them to
point to /World/Looks/... instead.

Usage:
    python migrate/fix_test_tube_materials.py
"""
from pxr import Usd, UsdShade

USD_PATH = "workflows/simbox/example_assets/task/pick_test_tube/test_tube/Aligned_obj.usd"

# Mapping: mesh prim path -> material path in new structure
MATERIAL_BINDINGS = {
    "/World/Aligned/_______005": "/World/Looks/SimPBR_Translucent",
    "/World/Aligned/tags/_______007": "/World/Looks/OmniPBR",
    "/World/Aligned/Test_Tube_lid/_______006": "/World/Looks/OmniPBR_01",
}


def fix():
    stage = Usd.Stage.Open(USD_PATH)

    for mesh_path, mat_path in MATERIAL_BINDINGS.items():
        mesh_prim = stage.GetPrimAtPath(mesh_path)
        mat_prim = stage.GetPrimAtPath(mat_path)

        if not mesh_prim.IsValid():
            print(f"  SKIP: {mesh_path} not found")
            continue
        if not mat_prim.IsValid():
            print(f"  SKIP: material {mat_path} not found")
            continue

        mat = UsdShade.Material(mat_prim)
        UsdShade.MaterialBindingAPI.Apply(mesh_prim)
        UsdShade.MaterialBindingAPI(mesh_prim).Bind(mat)
        print(f"  Bound {mesh_path} -> {mat_path}")

    stage.GetRootLayer().Save()
    print(f"\nSaved: {USD_PATH}")


if __name__ == "__main__":
    fix()