import asyncio
import logging
from typing import Callable, List, Optional

import grpc
import grpc.aio


from protos.BitAgentWorker_pb2_grpc import ConnectServiceStub

logger = logging.getLogger(__name__)


class GrpcClient:
    GRPC_CHANNEL_OPTIONS = [
        ("grpc.max_send_message_length", 10*1024*1024),
        ("grpc.max_receive_message_length", 10*1024*1024),
        # ("grpc.keep_alive_time", settings.GRPC_HEART_INTERVAL_SECONDS),
        # ("grpc.keep_alive_timeout", settings.GRPC_HEART_INTERVAL_SECONDS),
        # ("grpc.enable_retries", 1),
    ]
    connect_timeout_seconds: int = 60

    def __init__(self):
        self._address: str = ""
        self._use_ssl: bool = False
        self._header: List[tuple] = []
        self._is_disconnected: bool = False
        self._is_reconnecting: bool = False
        self._reconnect_lock = asyncio.Lock()
        self._reconnect_event = asyncio.Event()
        self._channel: Optional[grpc.aio.Channel] = None
        self._stub: Optional[ConnectServiceStub] = None
        self.state_changed_handlers: List[Callable[[bool], None]] = []

    @property
    def state(self) -> grpc.ChannelConnectivity:
        """获取通道状态"""
        if not self._channel:
            return grpc.ChannelConnectivity.SHUTDOWN
        return self._channel.get_state(try_to_connect=False)

    def is_channel_healthy(self) -> bool:
        """检查通道状态是否健康"""
        state = self.state
        return state in (
            grpc.ChannelConnectivity.CONNECTING,
            grpc.ChannelConnectivity.READY,
        )

    def set_header(self, token: str, client_id: str) -> None:
        """刷新token"""
        self._header = [("token", token), ("x-worker-id", client_id)]

    async def connect(self, address: str, use_ssl: bool) -> bool:
        """连接服务器"""
        if self._header.count == 0:
            raise ValueError("token and client_id must be set")

        if address is None or address == "":
            raise ValueError("grpc address must be set")

        logger.info(
            f"connect grpc to {address} use_ssl:{use_ssl} header:{str(self._header)}"
        )
        self._is_disconnected = False
        await self.shutdown()

        self._address = address
        self._use_ssl = use_ssl

        return await self.reconnect()

    async def disconnect(self) -> None:
        """断开连接"""
        if self._is_disconnected:
            return
        self._is_disconnected = True
        await self.shutdown()

    def create_all_stream(self) -> grpc.aio._call.StreamStreamCall:
        """创建双向流"""
        stub = ConnectServiceStub(self._channel)
        return stub.ReversalCall(metadata=self._header)

    async def _handle_rpc_exception(self, ex: grpc.aio.AioRpcError) -> bool:
        """处理RPC异常"""
        handled = False
        if ex.code() in (
            grpc.StatusCode.CANCELLED,
            grpc.StatusCode.UNAVAILABLE,
            grpc.StatusCode.DEADLINE_EXCEEDED,
            grpc.StatusCode.UNKNOWN,
        ):
            handled = await self.reconnect(force=True)
        return handled

    async def wait_reconnecting(self, timeout_seconds: int) -> bool:
        """等待重连"""
        if self._is_reconnecting:
            return await asyncio.wait_for(
                self._reconnect_event.wait(), timeout=timeout_seconds
            )
        return True

    def check_can_send(self, log: str) -> bool:
        """检查是否可以发送数据"""
        if self._is_reconnecting:
            logger.error(f"send error:{log} error，正在重连 {self.state}")
            return False
        if not self.is_channel_healthy():
            logger.error(f"send error:{log} error，连接失败 {self.state}")
            return False
        return True

    async def reconnect(self, force: bool = False) -> bool:
        """重新连接"""
        async with self._reconnect_lock:
            try:
                self._is_reconnecting = True
                self._reconnect_event.clear()
                if self._is_disconnected:
                    return False

                if force:
                    await self.shutdown()
                elif self.is_channel_healthy():
                    return True

                # 创建通道
                if not self._channel:
                    logger.info("create channel")
                    if self._use_ssl:
                        credits: grpc.ChannelCredentials = (
                            grpc.ssl_channel_credentials()
                        )
                        self._channel = grpc.aio.secure_channel(
                            self._address, credits, options=self.GRPC_CHANNEL_OPTIONS
                        )
                    else:
                        self._channel = grpc.aio.insecure_channel(
                            self._address, options=self.GRPC_CHANNEL_OPTIONS
                        )

                # 等待连接
                logger.info(f"wait channel ready to {self._address}")
                await asyncio.wait_for(
                    self._channel.channel_ready(), timeout=self.connect_timeout_seconds
                )

                logger.info("channel connected")
                # self.create_all_stream()  # 主要用于验证token
                # logger.info("create stream success")
                self.invoke_state_changed(True)
                return True

            except Exception as ex:
                if isinstance(ex, grpc.aio.AioRpcError):
                    logger.error(f"channel Connect error:{ex.code()} {ex.details()}")
                else:
                    logger.error(f"channel Connect error: {type(ex)}")
                    logger.exception(ex)
                await self.shutdown()
                await asyncio.sleep(3)  # 避免死循环重连

            finally:
                self._is_reconnecting = False
                self._reconnect_event.set()

            logger.info("channel error")
            return False

    async def shutdown(self) -> None:
        """关闭连接"""
        if self._channel:
            try:
                logger.info("channel Dispose")
                await asyncio.wait_for(
                    self._channel.close(), timeout=self.connect_timeout_seconds
                )
                logger.info("channel shutdown 1")
            except Exception as ex:
                logger.info(f"channel Shutdown error: {ex}")
            self._channel = None
            self.invoke_state_changed(False)

    # async def report_task_status(
    #     self, data: ReversalRequest, logger_msg: str = ""
    # ) -> bool:
    #     """上报任务状态"""
    #     log = logger_msg if logger_msg else data.jsonData
    #     try:
    #         if self._is_reconnecting:
    #             logger.error(f"send error:{log}, 正在重连")
    #             return False
    #         if not self.is_channel_healthy() and not await self.reconnect():
    #             logger.error(f"send error:{log}, 连接失败")
    #             return False

    #         stub = ConnectServiceStub(self._channel)
    #         response = await stub.reportTaskStatus(data, metadata=self._header)
    #         logger.info(f"{response.messageType} {response.jsonData} {log}")
    #         return True
    #     except Exception as e:
    #         if isinstance(e, grpc.aio.AioRpcError):
    #             logger.error(
    #                 f"send error RpcException:{log},{e.details()},{self.state}"
    #             )
    #             _ = self._handle_rpc_exception(e)
    #         else:
    #             logger.error(f"send error:{log},{e},{self.state}")

    #     return False

    # async def report_BitAgentWorker_log(
    #     self, data: ReversalRequest, can_retry: bool = True
    # ) -> bool:
    #     """上报机器人日志"""
    #     try:
    #         if self._is_reconnecting:
    #             logger.error("ReportRobotLog error，正在重连 " + str(self.state))
    #             return False

    #         if not self.is_channel_healthy() and not await self.reconnect():
    #             logger.error("ReportRobotLog error，连接失败 " + str(self.state))
    #             return False

    #         stub = ConnectServiceStub(self._channel)
    #         async with stub.reportRobotLog(metadata=self._header) as stream:
    #             await stream.write(data)
    #             await stream.done_writing()
    #             response = await stream.read()

    #             if not response.jsonData or "Fail" in response.jsonData:
    #                 return False

    #             return True

    #     except Exception as e:
    #         if isinstance(e, grpc.aio.AioRpcError):
    #             if can_retry and e.code() == grpc.StatusCode.UNKNOWN:
    #                 return await self.report_BitAgentWorker_log(data, False)
    #             logger.error(f"reportRobotLog error: {e.details()}")
    #             _ = self._handle_rpc_exception(e)
    #         else:
    #             logger.error(f"reportRobotLog error: {e}")

    #     return False

    def invoke_state_changed(self, state: bool) -> None:
        """状态改变"""
        for handler in self.state_changed_handlers:
            handler(state)
