import time
from typing import Callable, Optional

from google.protobuf.timestamp_pb2 import Timestamp

from clients._path_setup import ensure_paths

ensure_paths()

from clients.base_client import RobotClientBase
from cmvr.api import camera_command_pb2 as pb
from cmvr.api import camera_service_pb2_grpc as rpc
from cmvr.api import common_pb2


class CameraClient(RobotClientBase):
    """Client for CameraService RGB image RPCs."""

    def _init_stubs(self, channel):
        self.stub = rpc.CameraServiceStub(channel)

    @staticmethod
    def _make_header(device_id: str):
        header = common_pb2.CommandHeader.Request()
        header.device_id = device_id
        timestamp = Timestamp()
        timestamp.GetCurrentTime()
        header.timestamp.CopyFrom(timestamp)
        return header

    def start_camera(self, device_id: str):
        request = pb.StartCameraCommand.Request(header=self._make_header(device_id))
        return self.stub.StartCamera(request, timeout=self.timeout)

    def stop_camera(self, device_id: str):
        request = pb.StopCameraCommand.Request(header=self._make_header(device_id))
        return self.stub.StopCamera(request, timeout=self.timeout)

    def get_rgb_image(self, device_id: str):
        request = pb.GetRGBImageCommand.Request(header=self._make_header(device_id))
        return self.stub.GetRGBImage(request, timeout=self.timeout)

    def get_rgb_image_stream(
        self,
        device_id: str,
        stop_requested: Optional[Callable[[], bool]] = None,
        request_interval: float = 0.001,
    ):
        should_stop = stop_requested or (lambda: False)

        def request_gen():
            while not should_stop():
                yield pb.GetRGBImageStreamCommand.Request(
                    header=self._make_header(device_id),
                    eof=False,
                )
                time.sleep(request_interval)

            yield pb.GetRGBImageStreamCommand.Request(
                header=self._make_header(device_id),
                eof=True,
            )

        return self.stub.GetRGBImageStream(request_gen())


GetRGBImageStreamClient = CameraClient