"""
Repackage test tube USD to match InternDataEngine's expected format.

Target structure (same as Aligned_obj.usd):
  /World (Xform, defaultPrim, NO physics schemas)
    /Looks (Scope) - materials referenced from source
    /Aligned (Xform) - PhysicsRigidBodyAPI, PhysicsMassAPI
      /mesh (first child mesh, PhysicsCollisionAPI)
    /PhysicsMaterial (Material, PhysicsMaterialAPI)

Strategy:
  - /World is a clean Xform (no reference, no inherited schemas)
  - /World/Aligned references /Test_Tube_AA_01/Test_Tube from source
  - /World/Looks references /Test_Tube_AA_01/Looks from source
  - Physics applied as overrides on /World/Aligned

Usage:
    python migrate/repackage_test_tube.py
"""
from pxr import Usd, UsdGeom, UsdPhysics, UsdShade, Sdf, Gf
import os

SRC_PATH = os.path.abspath(
    "/home/tangger/LYT/maic_usd_assets_moudle/laboratory_equipment/Test_Tube/Test_Tube_AA_01.usdc"
)
DST_DIR = "workflows/simbox/example_assets/task/pick_test_tube/test_tube"
DST_PATH = os.path.join(DST_DIR, "Aligned_obj.usd")


def repackage():
    # Read source to understand structure
    src_stage = Usd.Stage.Open(SRC_PATH)
    src_dp = src_stage.GetDefaultPrim()
    print(f"Source: {SRC_PATH}")
    print(f"Source defaultPrim: {src_dp.GetPath()}")

    # Remove old output
    if os.path.exists(DST_PATH):
        os.remove(DST_PATH)

    # Create new stage
    dst_stage = Usd.Stage.CreateNew(DST_PATH)
    UsdGeom.SetStageMetersPerUnit(dst_stage, 1.0)
    UsdGeom.SetStageUpAxis(dst_stage, UsdGeom.Tokens.z)

    # --- /World: clean Xform, no reference, no inherited schemas ---
    world_xform = UsdGeom.Xform.Define(dst_stage, "/World")
    dst_stage.SetDefaultPrim(world_xform.GetPrim())

    # --- /World/Aligned: reference Test_Tube child from source ---
    aligned_xform = UsdGeom.Xform.Define(dst_stage, "/World/Aligned")
    aligned_prim = aligned_xform.GetPrim()
    aligned_prim.GetReferences().AddReference(
        SRC_PATH,
        f"{src_dp.GetPath()}/Test_Tube"
    )
    # Add RigidBody physics
    UsdPhysics.RigidBodyAPI.Apply(aligned_prim)
    mass_api = UsdPhysics.MassAPI.Apply(aligned_prim)
    mass_api.GetMassAttr().Set(0.005)
    print("Created /World/Aligned with RigidBodyAPI + MassAPI (0.005 kg)")

    # Set collision approximation on mesh children
    for child_prim in aligned_prim.GetAllChildren():
        if child_prim.HasAPI(UsdPhysics.CollisionAPI):
            mesh_api = UsdPhysics.MeshCollisionAPI(child_prim)
            if mesh_api:
                mesh_api.GetApproximationAttr().Set("convexHull")
                print(f"  Set convexHull on {child_prim.GetPath()}")

    # --- /World/Looks: reference Looks from source ---
    looks_prim = dst_stage.DefinePrim("/World/Looks", "Scope")
    looks_prim.GetReferences().AddReference(
        SRC_PATH,
        f"{src_dp.GetPath()}/Looks"
    )
    print("Created /World/Looks referencing source materials")

    # --- Fix material bindings: remap from source paths to new paths ---
    # Source materials are at /Test_Tube_AA_01/Looks/XXX
    # In our new stage they are at /World/Looks/XXX
    # The mesh bindings from the reference point to /Test_Tube_AA_01/Looks/XXX
    # which is outside scope. We need to override the bindings.
    src_material_map = {
        "SimPBR_Translucent": "/World/Looks/SimPBR_Translucent",
        "OmniPBR": "/World/Looks/OmniPBR",
        "OmniPBR_01": "/World/Looks/OmniPBR_01",
    }

    # Find all mesh prims under /World/Aligned and rebind materials
    for prim in dst_stage.Traverse():
        if not str(prim.GetPath()).startswith("/World/Aligned"):
            continue
        binding = UsdShade.MaterialBindingAPI(prim)
        if not binding:
            continue
        # Check if there's a direct binding
        mat_path_rel = binding.GetDirectBinding().GetMaterialPath()
        if mat_path_rel and str(mat_path_rel) != "":
            mat_name = str(mat_path_rel).split("/")[-1]
            if mat_name in src_material_map:
                new_mat_path = src_material_map[mat_name]
                new_mat = UsdShade.Material(dst_stage.GetPrimAtPath(new_mat_path))
                if new_mat:
                    UsdShade.MaterialBindingAPI.Apply(prim)
                    UsdShade.MaterialBindingAPI(prim).Bind(new_mat)
                    print(f"  Rebound material on {prim.GetPath()} -> {new_mat_path}")

    # --- /World/PhysicsMaterial ---
    phys_mat_prim = dst_stage.DefinePrim("/World/PhysicsMaterial", "Material")
    UsdPhysics.MaterialAPI.Apply(phys_mat_prim)
    phys_mat = UsdPhysics.MaterialAPI(phys_mat_prim)
    phys_mat.GetStaticFrictionAttr().Set(0.5)
    phys_mat.GetDynamicFrictionAttr().Set(0.5)
    phys_mat.GetRestitutionAttr().Set(0.1)
    print("Created /World/PhysicsMaterial")

    # Save
    dst_stage.GetRootLayer().Save()
    print(f"\nSaved: {DST_PATH}")

    # --- Verify ---
    print(f"\n{'='*60}")
    print("Verification:")
    print(f"{'='*60}")
    v_stage = Usd.Stage.Open(DST_PATH)
    dp = v_stage.GetDefaultPrim()
    print(f"DefaultPrim: {dp.GetPath()}")
    print(f"  Type: {dp.GetTypeName()}")
    print(f"  Schemas: {dp.GetAppliedSchemas()}")
    for c in dp.GetChildren():
        print(f"  /{c.GetName()} [{c.GetTypeName()}] schemas={c.GetAppliedSchemas()}")
        for gc in c.GetAllChildren():
            schemas = gc.GetAppliedSchemas()
            s = f" schemas={schemas}" if schemas else ""
            print(f"    /{gc.GetName()} [{gc.GetTypeName()}]{s}")


if __name__ == "__main__":
    repackage()