import time

import grpc

from clients._path_setup import ensure_paths

ensure_paths()

from cmvr.api import dexhand_service_pb2_grpc
from cmvr.api import arm_service_pb2_grpc


class RobotClientBase:
    def __init__(self, address="192.168.0.222:50052", timeout=2, retries=1):
        self.address = address
        self.timeout = timeout
        self.retries = retries
        self.channel = None
        self.stub = None
        self.stub_hand = None
        print(f"Attempting to connect to gRPC server: {self.address}")
        self._connect_with_retry()

    def _init_stubs(self, channel):
        self.stub = arm_service_pb2_grpc.ArmServiceStub(channel)
        self.stub_hand = dexhand_service_pb2_grpc.DexHandServiceStub(channel)

    def _connect_with_retry(self):
        attempt = 0
        while attempt <= self.retries:
            print(f"Connecting... (Attempt {attempt + 1})")
            try:
                self.channel = grpc.insecure_channel(self.address)
                grpc.channel_ready_future(self.channel).result(timeout=self.timeout)
                self._init_stubs(self.channel)
                print(f"Successfully connected to gRPC server: {self.address}")
                return
            except grpc.FutureTimeoutError:
                attempt += 1
                print(f"Connection timed out, attempt {attempt} failed")
                if attempt > self.retries:
                    raise ConnectionError(
                        f"Failed to connect to gRPC server {self.address} after {self.retries} retries."
                    )
                print(f"Waiting {self.timeout}s before retrying...")
                time.sleep(self.timeout)

    def close(self):
        if self.channel:
            self.channel.close()
            print(f"Closed connection to gRPC server {self.address}")