import grpc
from typing import Dict

from .enums import CMVRErrorCode
from .biohead_client import BioHeadClient


class CMVRGrpcClient:
    """CMVR gRPC 主客户端"""

    def __init__(self, server_address: str):
        self.server_address = server_address
        self.connected = False

        # 创建 gRPC 通道
        self.channel = grpc.insecure_channel(server_address)

        # 初始化各服务的存根
        self.generated = None  # 推迟导入
        self.biohead_stub = None
        self.biohead_map: Dict[str, BioHeadClient] = {}

        # 检查连接状态
        try:
            grpc.channel_ready_future(self.channel).result(timeout=5)
            self.connected = True
            self._import_generated()  # 在连接成功后才导入 generated
            self.biohead_stub = self.generated.biohead_service_pb2_grpc.BioHeadServiceStub(self.channel)
        except grpc.FutureTimeoutError:
            print(f"连接服务器超时: {server_address}")
            self.connected = False

    def _import_generated(self):
        """推迟导入 generated"""
        if self.generated is None:
            try:
                # 从正确的路径导入生成的模块
                from generated.cmvr.api import biohead_service_pb2_grpc
                from generated.cmvr.api import common_pb2, biohead_command_pb2
                self.generated = type('GeneratedModules', (), {
                    'biohead_service_pb2_grpc': biohead_service_pb2_grpc,
                    'common_pb2': common_pb2,
                    'biohead_command_pb2': biohead_command_pb2
                })
            except ImportError as e:
                print(f"导入生成的模块失败: {e}")
                print("请确保已生成 protobuf 代码")
                raise
        return self.generated

    def is_connected(self) -> bool:
        """检查连接状态"""
        return self.connected

    def get_biohead(self, device_id: str) -> BioHeadClient:
        """获取仿生头客户端"""
        if device_id not in self.biohead_map:
            self.biohead_map[device_id] = BioHeadClient(device_id, self.biohead_stub)
        return self.biohead_map[device_id]

    def close(self):
        """关闭所有连接"""
        # 关闭所有流式连接
        for biohead in self.biohead_map.values():
            biohead.close()

        # 关闭 gRPC 通道
        self.channel.close()

        # 清空所有映射
        self.biohead_map.clear()

        self.connected = False