from antioch.session.session import Session, SessionContainer
from common.message import Pose, RadarScan
from common.session.views.radar import AddRadar, GetRadar, GetRadarResponse, GetRadarScan, RadarConfig


class Radar(SessionContainer):
    """
    Radar view for time-synchronized scan data.

    Example:
        scene = Scene()
        radar = scene.get_radar(name="my_ark/my_module/my_radar")
        scan = radar.get_scan()
    """

    def __init__(self, path: str):
        """
        Initialize Radar view.

        :param path: USD path for the radar.
        """

        super().__init__()

        self._path = self._session.query_sim_rpc(
            endpoint="get_radar",
            payload=GetRadar(path=path),
            response_type=GetRadarResponse,
        ).path

    @classmethod
    def add(
        cls,
        path: str,
        config: RadarConfig,
        world_pose: Pose | None,
        local_pose: Pose | None,
    ) -> "Radar":
        """
        Add radar to the scene.

        :param path: USD path for the radar.
        :param config: Radar configuration.
        :param world_pose: Optional world pose.
        :param local_pose: Optional local pose.
        :return: The radar instance.
        """

        Session.get_current().query_sim_rpc(
            endpoint="add_radar",
            payload=AddRadar(
                path=path,
                config=config,
                world_pose=world_pose,
                local_pose=local_pose,
            ),
        )
        return cls(path)

    def get_scan(self) -> RadarScan | None:
        """
        Get radar scan data.

        :return: Radar scan with detections, or None if scan data is not ready.
        """

        scan = self._session.query_sim_rpc(
            endpoint="get_radar_scan",
            payload=GetRadarScan(path=self._path),
            response_type=RadarScan,
        )

        return scan
