import asyncio
import json
import logging
from collections import deque
from typing import Callable, Coroutine, Optional

import grpc.aio
from result import ResultModel


from grpc_client import GrpcClient
from protos.BitAgentWorker_pb2 import ReversalRequest, ReversalResponse

logger = logging.getLogger(__name__)


class GrpcListener:
    def __init__(self, client: GrpcClient):
        self._client: GrpcClient = client
        self.is_running: bool = False
        self._call_stream: Optional[grpc.aio._call.StreamStreamCall] = None
        self._stream_lock = asyncio.Lock()
        self._listen_task: Optional[asyncio.Task] = None
        self._request_id_set = set()
        self._request_id_deque = deque()
        self._ping_count = 0

        # 事件处理器
        self.message_received: (
            Optional[Callable[[str, str], Coroutine[None, None, ResultModel]]]
        ) = None
        client.state_changed_handlers.append(self._on_channel_state_changed)

    def _on_channel_state_changed(self, connect: bool) -> None:
        """处理通道状态变更"""
        if not connect:
            _ = self._close_call_stream()

    @property
    async def call_stream(self) -> Optional[grpc.aio._call.StreamStreamCall]:
        """获取调用流"""
        async with self._stream_lock:
            if not self.is_running:
                raise RuntimeError(f"{self.__class__.__name__} is Stop")
            if not self._call_stream:
                # 创建流
                if self._client.state not in (grpc.ChannelConnectivity.READY,):
                    return None
                stream = self._client.create_all_stream()
                logger.info("create listen stream success")
                self._client.invoke_state_changed(True)
                self._call_stream = stream
            return self._call_stream

    def start(self) -> None:
        """启动监听"""
        if self.is_running:
            return

        self.is_running = True
        self._listen_task = asyncio.create_task(self._listen_loop())

    async def _listen_loop(self) -> None:
        read_pass_count = 0
        logger.info("ListenHandle start")

        while self.is_running:
            try:
                stream = await self.call_stream
                if stream:
                    try:
                        response = await asyncio.wait_for(stream.read(), timeout=10)
                    except asyncio.TimeoutError:
                        logger.warning("grpc read timeout")
                        continue
                    if response:
                        read_pass_count = 0
                        # 1.ack
                        if response.body.startswith("#pong"):
                            if self._ping_count == 1:
                                logger.info(f"{response.body}")
                            continue

                        await stream.write(
                            ReversalResponse(requestId=response.id, type=0)
                        )
                        # 2.处理消息
                        asyncio.create_task(self._on_message_received(response))
                else:
                    # 如果网络不正确，直接触发重连
                    if not self._client.is_channel_healthy():
                        raise grpc.aio.AioRpcError(
                            grpc.StatusCode.CANCELLED,
                            grpc.aio.Metadata(),
                            grpc.aio.Metadata(),
                            details=f"state is {self._client.state}",
                        )
                    # 每2s记录一次日志
                    if read_pass_count % 40 == 0:
                        read_pass_count = 0
                        logger.error(
                            f"ListenHandle State error: {self._client.state} {stream is None}"
                        )
                        self._client.invoke_state_changed(False)

                    read_pass_count += 1
                    await asyncio.sleep(0.05)
            except asyncio.CancelledError:
                await asyncio.sleep(1)
                continue
            except Exception as e:
                logger.debug(f"ListenHandle exception:{type(e)}")
                if not self.is_running:
                    break
                if isinstance(e, grpc.aio.AioRpcError):
                    await self._close_call_stream()
                    self._client.invoke_state_changed(False)
                    logger.error(f"ListenHandle RpcError: {e.details()}")
                    if not await self._client._handle_rpc_exception(e):
                        if read_pass_count > 0:
                            await asyncio.sleep(1)
                        read_pass_count += 1
                else:
                    if read_pass_count > 0:
                        logger.error(f"ListenHandle exception: {e}")
                        await asyncio.sleep(1)
                    read_pass_count += 1

        logger.info("ListenHandle stop")

        await self.stop(None)

    async def _on_message_received(self, data: ReversalRequest) -> None:
        """处理接收到的消息"""
        try:
            # 1. 记录指令
            logger.info(f"收到指令：\n{data}")

            # 2. 指令判重
            if data.id in self._request_id_set:
                logger.warning(f"忽略重复指令：{data.id}")
                return
            self._request_id_set.add(data.id)
            self._request_id_deque.append(data.id)
            if len(self._request_id_set) > 100000:
                self._request_id_set.remove(
                    self._request_id_deque.popleft()
                )  # 最早的指令移除

            # 3. 消费指令
            consume_result: ResultModel = ResultModel(code=500, message="client error")
            try:
                if self.message_received:
                    parts = data.body.split("\n", maxsplit=2)
                    if len(parts) == 3:
                        request_path = parts[0].strip("/")
                        request_data = parts[2]
                        consume_result = await self.message_received(
                            request_path, request_data
                        )
                    else:
                        consume_result = ResultModel(code=500, message="指令格式错误")
                        logger.error(f"指令格式错误：{data.body}")

            except Exception as e:
                consume_result = ResultModel(code=500, message=str(e))
                logger.exception(f"消费指令异常:{data} {e}")

            # 3. 回复执行结果
            responseBody = f"{consume_result.code}\n{consume_result.message}"
            if consume_result.code == 200 and consume_result.data is not None:
                try:
                    response_headers = json.dumps(
                        {},
                        default=lambda o: o.__dict__,
                        ensure_ascii=False,
                    )
                    response_data = json.dumps(
                        consume_result.data,
                        default=lambda o: o.__dict__,
                        ensure_ascii=False,
                    )
                    responseBody = (
                        f"{consume_result.code}\n{response_headers}\n{response_data}"
                    )
                except Exception as e:
                    logger.exception(f"json encode error:{consume_result.data} {e}")
            await self.send_to_all_stream(data.id, responseBody, retry_count=3)
        except Exception as e:
            logger.exception(f"on_message_received error:{e}")

    async def stop(self, data: Optional[ReversalResponse] = None) -> None:
        """停止监听"""
        if data and self._call_stream:
            try:
                await self._call_stream.write(data)
                logger.info("SendToStream DisConnect success")
            except Exception as ex:
                logger.error(f"send error:SendToStream DisConnect,{ex}")

        self.is_running = False
        if self._listen_task:
            self._listen_task.cancel()
            self._listen_task = None

        await self._close_call_stream()

    async def _close_call_stream(self) -> None:
        """关闭调用流"""
        async with self._stream_lock:
            if self._call_stream:
                try:
                    await self._call_stream.done_writing()
                except Exception:
                    pass
                try:
                    self._call_stream.cancel()
                except Exception:
                    logger.exception("Stream dispose error")
                self._call_stream = None

    async def send_to_all_stream(
        self, requestId: str, body: str, type: int = 1, retry_count: int = 0
    ) -> bool:
        """发送数据到所有流"""
        if type == -1:
            if self._ping_count == 6:
                self._ping_count = 0
            self._ping_count += 1
        if len(body) < 2000:
            log = f"{requestId} content:{body}"
        else:
            log = f"{requestId} len:{len(body)} content:{body[:2000]}..."
        count = retry_count
        while count >= 0:
            try:
                # 判断是否可以发送数据
                if retry_count > 0:
                    if count < retry_count:
                        logger.info("send retry ")
                    await self._client.wait_reconnecting(
                        30
                    )
                else:
                    if not self._client.check_can_send(log):
                        await self._close_call_stream()
                        continue

                # 分段发送数据
                partitionIndex = 0
                stream = await self.call_stream
                if not stream:
                    logger.error(f"call_stream is null:{log}")
                    continue
                for i in range(0, len(body), 1000):
                    if partitionIndex > 0:
                        logger.info(f"partition:{partitionIndex}")
                    data = ReversalResponse(
                        requestId=requestId,
                        type=type,
                        partitionIndex=partitionIndex,
                        partitionData=body[i : i + 1000],
                    )
                    await stream.write(data)
                    partitionIndex += 1
                # 发送结束标志
                if type == 1:
                    await stream.write(
                        ReversalResponse(requestId=data.requestId, type=2)
                    )
                if type != -1 or self._ping_count == 1:
                    logger.info(f"send {log}")
                return True
            except Exception as e:
                if isinstance(e, grpc.aio.AioRpcError):
                    logger.error(f"send RpcException:{log},{str(e)},{e.details()}")
                    await self._client._handle_rpc_exception(e)
                elif isinstance(e, asyncio.InvalidStateError):
                    logger.error(f"send StateError:{log},{str(e)},{self._client.state}")
                    await self._close_call_stream()
                    await self._client.reconnect(force=True)
                else:
                    logger.error(f"send Exception:{log},{str(e)},{self._client.state}")
            finally:
                count -= 1
        return False
