# -*- coding: utf-8 -*-

__all__ = ["BLiveClient"]

import json
import zlib
import struct
import asyncio
import logging
import ssl as ssl_
from enum import IntEnum

# code from xfgryujk
from abc import abstractmethod
from collections import namedtuple
from typing import *  # ??这是什么粗暴的import方式

import aiohttp

logger = logging.getLogger(__name__)


ROOM_INIT_URL = "https://api.live.bilibili.com/xlive/web-room/v1/index/getInfoByRoom"
DANMAKU_SERVER_CONF_URL = (
    "https://api.live.bilibili.com/xlive/web-room/v1/index/getDanmuInfo"
)
DEFAULT_DANMAKU_SERVER_LIST = [
    {
        "host": "broadcastlv.chat.bilibili.com",
        "port": 2243,
        "wss_port": 443,
        "ws_port": 2244,
    }
]

HEADER_STRUCT = struct.Struct(">I2H2I")
HeaderTuple = namedtuple(
    "HeaderTuple", ("pack_len", "raw_header_size", "ver", "operation", "seq_id")
)
WS_BODY_PROTOCOL_VERSION_INFLATE = 0
WS_BODY_PROTOCOL_VERSION_NORMAL = 1
WS_BODY_PROTOCOL_VERSION_DEFLATE = 2


# go-common\app\service\main\broadcast\model\operation.go
class Operation(IntEnum):
    HANDSHAKE = 0
    HANDSHAKE_REPLY = 1
    HEARTBEAT = 2
    HEARTBEAT_REPLY = 3
    SEND_MSG = 4
    SEND_MSG_REPLY = 5
    DISCONNECT_REPLY = 6
    AUTH = 7
    AUTH_REPLY = 8
    RAW = 9
    PROTO_READY = 10
    PROTO_FINISH = 11
    CHANGE_ROOM = 12
    CHANGE_ROOM_REPLY = 13
    REGISTER = 14
    REGISTER_REPLY = 15
    UNREGISTER = 16
    UNREGISTER_REPLY = 17
    # B站业务自定义OP
    # MinBusinessOp = 1000
    # MaxBusinessOp = 10000


class InitError(Exception):
    """初始化失败"""


class BLiveClient:
    def __init__(
        self,
        room_id,
        uid=0,
        session: aiohttp.ClientSession = None,
        heartbeat_interval=30,
        ssl=True,
        loop=None,
    ):
        """
        :param room_id: URL中的房间ID，可以为短ID
        :param uid: B站用户ID，0表示未登录
        :param session: cookie、连接池
        :param heartbeat_interval: 发送心跳包的间隔时间（秒）
        :param ssl: True表示用默认的SSLContext验证，False表示不验证，也可以传入SSLContext
        :param loop: 协程事件循环
        """
        # 用来init_room的临时房间ID
        self._tmp_room_id = room_id
        # 调用init_room后初始化
        self._room_id = self._room_short_id = self._room_owner_uid = None
        # [{host: "tx-bj4-live-comet-04.chat.bilibili.com", port: 2243, wss_port: 443, ws_port: 2244}, ...]
        self._host_server_list = None
        self._host_server_token = None
        self._uid = uid

        if loop is not None:
            self._loop = loop
        elif session is not None:
            # noinspection PyDeprecation
            self._loop = session.loop
        else:
            self._loop = asyncio.get_event_loop()
        self._future = None

        if session is None:
            self._session = aiohttp.ClientSession(
                loop=self._loop, timeout=aiohttp.ClientTimeout(total=10)
            )
            self._own_session = True
        else:
            self._session = session
            self._own_session = False
            # noinspection PyDeprecation
            if self._session.loop is not self._loop:
                raise RuntimeError("BLiveClient and session has to use same event loop")

        self._heartbeat_interval = heartbeat_interval
        # noinspection PyProtectedMember
        self._ssl = ssl if ssl else ssl_._create_unverified_context()
        self._websocket = None
        self._heartbeat_timer_handle = None

    @property
    def is_running(self):
        return self._future is not None

    @property
    def room_id(self):
        """
        房间ID，调用init_room后初始化
        """
        return self._room_id

    @property
    def room_short_id(self):
        """
        房间短ID，没有则为0，调用init_room后初始化
        """
        return self._room_short_id

    @property
    def room_owner_uid(self):
        """
        主播ID，调用init_room后初始化
        """
        return self._room_owner_uid

    async def close(self):
        """
        如果session是自己创建的则关闭session
        """
        if self._own_session:
            await self._session.close()

    def start(self):
        """
        创建相关的协程，不会执行事件循环
        :return: 协程的future
        """
        if self._future is not None:
            raise RuntimeError("This client is already running")
        self._future = asyncio.ensure_future(self._message_loop(), loop=self._loop)
        self._future.add_done_callback(self.__on_message_loop_done)
        return self._future

    def __on_message_loop_done(self, future):
        self._future = None
        logger.debug("room %s 消息协程结束", self.room_id)
        exception = future.exception()
        if exception is not None:
            logger.exception(
                "room %s 消息协程异常结束：",
                self.room_id,
                exc_info=(type(exception), exception, exception.__traceback__),
            )

    def stop(self):
        """
        停止相关的协程
        :return: 协程的future
        """
        if self._future is None:
            raise RuntimeError("This client is not running")
        self._future.cancel()
        return self._future

    async def init_room(self):
        """
        :return: True代表没有降级，如果需要降级后还可用，重载这个函数返回True
        """
        res = True
        if not await self._init_room_id_and_owner():
            res = False
            # 失败了则降级
            self._room_id = self._room_short_id = self._tmp_room_id
            self._room_owner_uid = 0

        if not await self._init_host_server():
            res = False
            # 失败了则降级
            self._host_server_list = DEFAULT_DANMAKU_SERVER_LIST
            self._host_server_token = None
        return res

    async def _init_room_id_and_owner(self):
        try:
            async with self._session.get(
                ROOM_INIT_URL, params={"room_id": self._tmp_room_id}, ssl=self._ssl
            ) as res:
                if res.status != 200:
                    logger.warning(
                        "room %d init_room失败：%d %s",
                        self._tmp_room_id,
                        res.status,
                        res.reason,
                    )
                    return False
                data = await res.json()
                if data["code"] != 0:
                    logger.warning(
                        "room %d init_room失败：%s", self._tmp_room_id, data["message"]
                    )
                    return False
                if not self._parse_room_init(data["data"]):
                    return False
        except (aiohttp.ClientConnectionError, asyncio.TimeoutError):
            logger.exception("room %d init_room失败：", self._tmp_room_id)
            return False
        return True

    def _parse_room_init(self, data):
        room_info = data["room_info"]
        self._room_id = room_info["room_id"]
        self._room_short_id = room_info["short_id"]
        self._room_owner_uid = room_info["uid"]
        return True

    async def _init_host_server(self):
        try:
            async with self._session.get(
                DANMAKU_SERVER_CONF_URL,
                params={"id": self._room_id, "type": 0},
                ssl=self._ssl,
            ) as res:
                if res.status != 200:
                    logger.warning(
                        "room %d getConf失败：%d %s", self._room_id, res.status, res.reason
                    )
                    return False
                data = await res.json()
                if data["code"] != 0:
                    logger.warning(
                        "room %d getConf失败：%s", self._room_id, data["message"]
                    )
                    return False
                if not self._parse_danmaku_server_conf(data["data"]):
                    return False
        except (aiohttp.ClientConnectionError, asyncio.TimeoutError):
            logger.exception("room %d getConf失败：", self._room_id)
            return False
        return True

    def _parse_danmaku_server_conf(self, data):
        self._host_server_list = data["host_list"]
        self._host_server_token = data["token"]
        if not self._host_server_list:
            logger.warning("room %d getConf失败：host_server_list为空", self._room_id)
            return False
        return True

    @staticmethod
    def _make_packet(data, operation):
        body = json.dumps(data).encode("utf-8")
        header = HEADER_STRUCT.pack(
            HEADER_STRUCT.size + len(body), HEADER_STRUCT.size, 1, operation, 1
        )
        return header + body

    async def _send_auth(self):
        auth_params = {
            "uid": self._uid,
            "roomid": self._room_id,
            "protover": 2,
            "platform": "web",
            "clientver": "1.14.3",
            "type": 2,
        }
        if self._host_server_token is not None:
            auth_params["key"] = self._host_server_token
        await self._websocket.send_bytes(self._make_packet(auth_params, Operation.AUTH))

    async def _message_loop(self):
        # 如果之前未初始化则初始化
        if self._host_server_token is None:
            if not await self.init_room():
                raise InitError("初始化失败")

        retry_count = 0
        while True:
            try:
                # 连接
                host_server = self._host_server_list[
                    retry_count % len(self._host_server_list)
                ]
                async with self._session.ws_connect(
                    f'wss://{host_server["host"]}:{host_server["wss_port"]}/sub',
                    receive_timeout=self._heartbeat_interval + 5,
                    ssl=self._ssl,
                ) as websocket:
                    self._websocket = websocket
                    await self._send_auth()
                    self._heartbeat_timer_handle = self._loop.call_later(
                        self._heartbeat_interval, self._on_send_heartbeat
                    )

                    # 处理消息
                    async for message in websocket:  # type: aiohttp.WSMessage
                        retry_count = 0
                        if message.type != aiohttp.WSMsgType.BINARY:
                            logger.warning(
                                "room %d 未知的websocket消息：type=%s %s",
                                self.room_id,
                                message.type,
                                message.data,
                            )
                            continue

                        try:
                            await self._handle_message(message.data)
                        except asyncio.CancelledError:
                            logger.warn(f"{self.room_id} 程序被強制取消。")
                            raise
                        except Exception:
                            logger.exception("room %d 处理消息时发生错误：", self.room_id)

            except asyncio.CancelledError:
                break
            except (aiohttp.ClientConnectionError, asyncio.TimeoutError):
                # 重连
                pass
            except ssl_.SSLError:
                logger.exception("SSL错误：")
                # 证书错误时无法重连
                break
            finally:
                self._websocket = None
                if self._heartbeat_timer_handle is not None:
                    self._heartbeat_timer_handle.cancel()
                    self._heartbeat_timer_handle = None

            retry_count += 1
            logger.warning("room %d 掉线重连中%d", self.room_id, retry_count)
            try:
                await asyncio.sleep(1)
            except asyncio.CancelledError:
                break

    def _on_send_heartbeat(self):
        coro = self._websocket.send_bytes(self._make_packet({}, Operation.HEARTBEAT))
        asyncio.ensure_future(coro, loop=self._loop)
        self._heartbeat_timer_handle = self._loop.call_later(
            self._heartbeat_interval, self._on_send_heartbeat
        )

    async def _handle_message(self, data):
        offset = 0
        while offset < len(data):
            try:
                header = HeaderTuple(*HEADER_STRUCT.unpack_from(data, offset))
            except struct.error:
                break

            if header.operation == Operation.HEARTBEAT_REPLY:
                popularity = int.from_bytes(
                    data[offset + HEADER_STRUCT.size : offset + HEADER_STRUCT.size + 4],
                    "big",
                )
                await self._on_receive_popularity(popularity)

            elif header.operation == Operation.SEND_MSG_REPLY:
                body = data[offset + HEADER_STRUCT.size : offset + header.pack_len]
                if header.ver == WS_BODY_PROTOCOL_VERSION_DEFLATE:
                    body = await self._loop.run_in_executor(None, zlib.decompress, body)
                    await self._handle_message(body)
                else:
                    try:
                        body = json.loads(body.decode("utf-8"))
                        await self._handle_command(body)
                    except Exception:
                        logger.error("body: %s", body)
                        raise

            elif header.operation == Operation.AUTH_REPLY:
                await self._websocket.send_bytes(
                    self._make_packet({}, Operation.HEARTBEAT)
                )

            else:
                body = data[offset + HEADER_STRUCT.size : offset + header.pack_len]
                logger.warning(
                    "room %d 未知包类型：operation=%d %s%s",
                    self.room_id,
                    header.operation,
                    header,
                    body,
                )

            offset += header.pack_len

    async def _handle_command(self, command):
        if isinstance(command, list):
            for one_command in command:
                await self._handle_command(one_command)
            return
        cmd = command.get("cmd", "")
        pos = cmd.find(":")  # 2019-5-29 B站弹幕升级新增了参数
        if pos != -1:
            cmd = cmd[:pos]
        await self.on_command_received(cmd, command)

    @abstractmethod
    async def on_command_received(self, cmd, data):
        pass

    @abstractmethod
    async def _on_receive_popularity(self, popularity: int):
        pass
