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

import grpc.aio

from models import ResultModel

from body import format_response_body, parse_request_body
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)
        # 新增：发送队列和发送任务
        self._send_queue = asyncio.Queue()
        self._send_task: Optional[asyncio.Task] = None

    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())
        # 启动发送任务
        if not self._send_task or self._send_task.done():
            self._send_task = asyncio.create_task(self._send_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.info(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) > 3000:
                self._request_id_set.remove(
                    self._request_id_deque.popleft()
                )  # 最早的指令移除

            # 3. 消费指令
            consume_result: ResultModel = ResultModel.error("client error")
            try:
                request_data = parse_request_body(data.body)
                if request_data is None:
                    consume_result = ResultModel.error("指令格式错误")
                elif self.message_received:
                    consume_result = await self.message_received(
                        request_data.path, request_data.data
                    )
            except Exception as e:
                consume_result = ResultModel.error(str(e))
                logger.exception(f"消费指令异常:{data} {e}")

            # 3. 回复执行结果
            log = consume_result.log_str()
            await self.send_to_all_stream(
                data.id, format_response_body(consume_result), retry_count=3, log=log
            )
        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

        # 停止发送任务
        if self._send_task:
            self._send_task.cancel()
            self._send_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_loop(self):
        while self.is_running:
            try:
                send_args = await self._send_queue.get()
                *real_args, fut = send_args
                try:
                    result = await self._do_send_to_all_stream(*real_args)
                except Exception as e:
                    logger.exception(f"Send loop error: {e}")
                    result = False
                finally:
                    if fut:
                        fut.set_result(result)
            except asyncio.CancelledError:
                break

    async def send_to_all_stream(
        self,
        requestId: str,
        body: str,
        type: int = 1,
        retry_count: int = 0,
        log: Optional[str] = None,
    ) -> bool:
        """入队，串行发送数据"""
        if type == -1:
            if not self._send_queue.empty():
                logger.debug("ignore:ping")
                return True
            if self._ping_count == 6:
                self._ping_count = 0
            self._ping_count += 1
            log = "#ping"
        else:
            if not log:
                log = f'"{requestId}":\n{body}'
            else:
                log = f'"{requestId}":\n{log}'

        fut = asyncio.get_event_loop().create_future()
        await self._send_queue.put((requestId, body, type, retry_count, log, fut))
        return await fut

    async def _do_send_to_all_stream(
        self,
        requestId: str,
        body: str,
        type: int = 1,
        retry_count: int = 0,
        log: Optional[str] = None,
    ) -> bool:
        try:
            count = retry_count
            while count >= 0:
                try:
                    # 判断是否可以发送数据
                    if retry_count > 0:
                        if count < retry_count:
                            logger.info("send retry ")
                        await self._client.wait_reconnecting(
                            60
                        )
                    else:
                        if not self._client.check_can_send(log or ""):
                            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), 1024*1024*10):
                        if partitionIndex > 0:
                            logger.info(f"partition:{partitionIndex}")
                        data = ReversalResponse(
                            requestId=requestId,
                            type=type,
                            partitionIndex=partitionIndex,
                            partitionData=body[
                                i : i + 1024*1024*10
                            ],
                        )
                        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
        except Exception as e:
            logger.exception(f"_do_send_to_all_stream error: {e}")
            return False
