from typing import List, Optional

import attr

from robogym.envs.rearrange.common.utils import (
    get_mesh_bounding_box,
    make_mesh_object,
    make_target,
)
from robogym.envs.rearrange.simulation.base import (
    ObjectGroupConfig,
    RearrangeSimParameters,
    RearrangeSimulationInterface,
)


@attr.s(auto_attribs=True)
class MeshObjectGroupConfig(ObjectGroupConfig):
    # Mesh files for constructing objects
    mesh_files: List[str] = ["placeholder.stl"]


@attr.s(auto_attribs=True)
class MeshRearrangeSimParameters(RearrangeSimParameters):
    # Overwrite the object group config type.
    object_groups: List[MeshObjectGroupConfig] = None  # type: ignore

    def __attrs_post_init__(self):
        if self.object_groups is None:
            self.object_groups = [
                MeshObjectGroupConfig(count=1, object_ids=[obj_id])
                for obj_id in range(self.num_objects)
            ]

    # Mesh files for each object. If not specified default placeholder
    # mesh will be used.
    mesh_files: Optional[List[List[str]]] = None

    # Global scale factor of mesh
    mesh_scale: float = 1.0


class MeshRearrangeSim(RearrangeSimulationInterface[MeshRearrangeSimParameters]):
    """
    Move around a object generated by customized meshes on the table.
    The meshes can be selected from YCB or ShapeNet dataset.
    """

    @classmethod
    def make_objects_xml(cls, xml, simulation_params: MeshRearrangeSimParameters):
        xmls = []
        for obj_group in simulation_params.object_groups:
            # Use a placeholder stl is no mesh files specified in parameters if no mesh files
            # are defined. This should only happen during environment creation.
            for i in obj_group.object_ids:
                object_xml = make_mesh_object(
                    f"object{i}",
                    obj_group.mesh_files,
                    scale=obj_group.scale * simulation_params.mesh_scale,
                )
                xmls.append((object_xml, make_target(object_xml)))

        assert len(xmls) == simulation_params.num_objects
        return xmls

    def _get_bounding_box(self, object_name):
        return get_mesh_bounding_box(self.mj_sim, object_name)
