from metasim.cfg.objects import RigidObjMetaCfg
from metasim.constants import PhysicStateType
from metasim.utils import configclass

from .rlbench_task_metacfg import RLBenchTaskMetaCfg


@configclass
class LightBulbInMetaCfg(RLBenchTaskMetaCfg):
    episode_length = 200
    traj_filepath = "roboverse_data/trajs/rlbench/light_bulb_in/v2"
    objects = [
        RigidObjMetaCfg(
            name="bulb0",
            filepath="roboverse_data/assets/rlbench/light_bulb_in/bulb0/usd/bulb0.usd",
            physics=PhysicStateType.RIGIDBODY,
        ),
        RigidObjMetaCfg(
            name="bulb1",
            filepath="roboverse_data/assets/rlbench/light_bulb_in/bulb1/usd/bulb1.usd",
            physics=PhysicStateType.RIGIDBODY,
        ),
        RigidObjMetaCfg(
            name="bulb_holder0",
            filepath="roboverse_data/assets/rlbench/light_bulb_in/bulb_holder0/usd/bulb_holder0.usd",
            physics=PhysicStateType.GEOM,
        ),
        RigidObjMetaCfg(
            name="bulb_holder1",
            filepath="roboverse_data/assets/rlbench/light_bulb_in/bulb_holder1/usd/bulb_holder1.usd",
            physics=PhysicStateType.GEOM,
        ),
        RigidObjMetaCfg(
            name="lamp_base",
            filepath="roboverse_data/assets/rlbench/light_bulb_in/lamp_base/usd/lamp_base.usd",
            physics=PhysicStateType.GEOM,
        ),
    ]
    # TODO: add checker
