import asyncio
import codecs
import gzip
import hashlib
import json
import random
import string
import subprocess
import threading
import time
import traceback
import urllib.parse
from contextlib import contextmanager
from unittest.mock import patch
from datetime import datetime
import httpx
import requests
import websocket
from loguru import logger
from py_mini_racer import MiniRacer

from protobuf.douyin import *
from message_processor import *
from helper import resource_path, PromptQueue, MessageType


@contextmanager
def patched_popen_encoding(encoding='utf-8'):
    original_popen_init = subprocess.Popen.__init__

    def new_popen_init(self, *args, **kwargs):
        kwargs['encoding'] = encoding
        original_popen_init(self, *args, **kwargs)

    with patch.object(subprocess.Popen, '__init__', new_popen_init):
        yield


def generateSignature(wss, script_file='sign.js'):
    """
    出现gbk编码问题则修改 python模块subprocess.py的源码中Popen类的__init__函数参数encoding值为 "utf-8"
    """
    params = ("live_id,aid,version_code,webcast_sdk_version,"
              "room_id,sub_room_id,sub_channel_id,did_rule,"
              "user_unique_id,device_platform,device_type,ac,"
              "identity").split(',')
    wss_params = urllib.parse.urlparse(wss).query.split('&')
    wss_maps = {i.split('=')[0]: i.split("=")[-1] for i in wss_params}
    tpl_params = [f"{i}={wss_maps.get(i, '')}" for i in params]
    param = ','.join(tpl_params)
    md5 = hashlib.md5()
    md5.update(param.encode())
    md5_param = md5.hexdigest()

    with codecs.open(script_file, 'r', encoding='utf8') as f:
        script = f.read()

    ctx = MiniRacer()
    ctx.eval(script)

    try:
        signature = ctx.call("get_sign", md5_param)
        return signature
    except Exception as e:
        logger.error(e)

    # 以下代码对应js脚本为sign_v0.js
    # context = execjs.compile(script)
    # with patched_popen_encoding(encoding='utf-8'):
    #     ret = context.call('getSign', {'X-MS-STUB': md5_param})
    # return ret.get('X-Bogus')


def generateMsToken(length=107):
    """
    产生请求头部cookie中的msToken字段，其实为随机的107位字符
    :param length:字符位数
    :return:msToken
    """
    random_str = ''
    base_str = string.ascii_letters + string.digits + '=_'
    _len = len(base_str) - 1
    for _ in range(length):
        random_str += base_str[random.randint(0, _len)]
    return random_str


class DouyinLiveWebFetcher:

    def __init__(self, ws_open_event, queue: PromptQueue):
        """
        直播间弹幕抓取对象
        :param live_id: 直播间的直播id，打开直播间web首页的链接如：https://live.douyin.com/261378947940，
                        其中的261378947940即是live_id
        """
        self.__ttwid = None
        self.__room_id = None
        self.live_id = LiveChatConfig().live_id
        logger.info(f'self.live_id: {self.live_id}')
        self.live_url = "https://live.douyin.com/"
        self.user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) " \
                          "Chrome/120.0.0.0 Safari/537.36"
        self.queue = queue
        self.ws_open_event = ws_open_event

    def start(self):
        try:
            self._connectWebSocket()
        except Exception as e:
            live_chat_config.update_chat_enable_status('启动失败')
            logger.error(traceback.format_exc())
            raise e

    def stop(self):
        self.ws.close()

    @property
    def ttwid(self):
        """
        产生请求头部cookie中的ttwid字段，访问抖音网页版直播间首页可以获取到响应cookie中的ttwid
        :return: ttwid
        """
        if self.__ttwid:
            return self.__ttwid
        headers = {
            "User-Agent": self.user_agent,
        }
        try:
            response = requests.get(self.live_url, headers=headers)
            response.raise_for_status()
        except Exception as err:
            logger.error("【X】Request the live url error: ", err)
        else:
            self.__ttwid = response.cookies.get('ttwid')
            return self.__ttwid

    @property
    def room_id(self):
        """
        根据直播间的地址获取到真正的直播间roomId，有时会有错误，可以重试请求解决
        :return:room_id
        """
        if self.__room_id:
            return self.__room_id
        url = self.live_url + self.live_id
        headers = {
            "User-Agent": self.user_agent,
            "cookie": f"ttwid={self.ttwid}&msToken={generateMsToken()}; __ac_nonce=0123407cc00a9e438deb4",
        }
        try:
            response = requests.get(url, headers=headers)
            response.raise_for_status()
        except Exception as err:
            logger.error("【X】Request the live room url error: ", err)
        else:
            match = re.search(r'roomId\\":\\"(\d+)\\"', response.text)
            if match is None or len(match.groups()) < 1:
                logger.error("【X】No match found for roomId")

            self.__room_id = match.group(1)

            return self.__room_id

    def get_room_status(self):
        """
        获取直播间开播状态:
        room_status: 2 直播已结束
        room_status: 0 直播进行中
        """
        url = ('https://live.douyin.com/webcast/room/web/enter/?aid=6383'
               '&app_name=douyin_web&live_id=1&device_platform=web&language=zh-CN&enter_from=web_live'
               '&cookie_enabled=true&screen_width=1536&screen_height=864&browser_language=zh-CN&browser_platform=Win32'
               '&browser_name=Edge&browser_version=133.0.0.0'
               f'&web_rid={self.live_id}'
               f'&room_id_str={self.room_id}'
               '&enter_source=&is_need_double_stream=false&insert_task_id=&live_reason='
               '&msToken=&a_bogus=')
        resp = requests.get(url, headers={
            'User-Agent': self.user_agent,
            'Cookie': f'ttwid={self.ttwid};'
        })
        data = resp.json().get('data')
        if data:
            room_status = data.get('room_status')
            user = data.get('user')
            user_id = user.get('id_str')
            nickname = user.get('nickname')
            logger.info(f"【{nickname}】[{user_id}]直播间：{['正在直播', '已结束'][bool(room_status)]}.")

    def _connectWebSocket(self):
        """
        连接抖音直播间websocket服务器，请求直播间数据
        """
        wss = ("wss://webcast100-ws-web-lq.douyin.com/webcast/im/push/v2/?app_name=douyin_web"
               "&version_code=180800&webcast_sdk_version=1.0.14-beta.0"
               "&update_version_code=1.0.14-beta.0&compress=gzip&device_platform=web&cookie_enabled=true"
               "&screen_width=1536&screen_height=864&browser_language=zh-CN&browser_platform=Win32"
               "&browser_name=Mozilla"
               "&browser_version=5.0%20(Windows%20NT%2010.0;%20Win64;%20x64)%20AppleWebKit/537.36%20(KHTML,"
               "%20like%20Gecko)%20Chrome/126.0.0.0%20Safari/537.36"
               "&browser_online=true&tz_name=Asia/Shanghai"
               "&cursor=d-1_u-1_fh-7392091211001140287_t-1721106114633_r-1"
               f"&internal_ext=internal_src:dim|wss_push_room_id:{self.room_id}|wss_push_did:7319483754668557238"
               f"|first_req_ms:1721106114541|fetch_time:1721106114633|seq:1|wss_info:0-1721106114633-0-0|"
               f"wrds_v:7392094459690748497"
               f"&host=https://live.douyin.com&aid=6383&live_id=1&did_rule=3&endpoint=live_pc&support_wrds=1"
               f"&user_unique_id=7319483754668557238&im_path=/webcast/im/fetch/&identity=audience"
               f"&need_persist_msg_count=15&insert_task_id=&live_reason=&room_id={self.room_id}&heartbeatDuration=0")

        script_file = resource_path('sign.js')
        signature = generateSignature(wss, script_file)
        wss += f"&signature={signature}"

        headers = {
            "cookie": f"ttwid={self.ttwid}",
            'user-agent': self.user_agent,
        }
        self.ws = websocket.WebSocketApp(wss,
                                         header=headers,
                                         on_open=self._wsOnOpen,
                                         on_message=self._wsOnMessage,
                                         on_error=self._wsOnError,
                                         on_close=self._wsOnClose)
        try:
            self.ws.run_forever()
        except Exception:
            self.stop()
            raise

    def _sendHeartbeat(self):
        """
        发送心跳包
        """
        while True:
            try:
                heartbeat = PushFrame(payload_type='hb').SerializeToString()
                self.ws.send(heartbeat, websocket.ABNF.OPCODE_PING)
                # print("【√】发送心跳包")
            except Exception as e:
                logger.error("【X】心跳包检测错误: ", e)
                break
            else:
                time.sleep(5)

    def _wsOnOpen(self, ws):
        """
        连接建立成功
        """
        logger.info("【√】WebSocket连接成功.")
        self.ws_open_event.set()
        threading.Thread(target=self._sendHeartbeat).start()

    def _wsOnMessage(self, ws, message):
        """
        接收到数据
        :param ws: websocket实例
        :param message: 数据
        """

        # 根据proto结构体解析对象
        package = PushFrame().parse(message)
        response = Response().parse(gzip.decompress(package.payload))

        # 返回直播间服务器链接存活确认消息，便于持续获取数据
        if response.need_ack:
            ack = PushFrame(log_id=package.log_id,
                            payload_type='ack',
                            payload=response.internal_ext.encode('utf-8')
                            ).SerializeToString()
            ws.send(ack, websocket.ABNF.OPCODE_BINARY)

        # 根据消息类别解析消息体
        for msg in response.messages_list:
            method = msg.method
            try:
                funcs = {
                    'WebcastChatMessage': parse_chat_msg,  # 聊天消息
                    'WebcastGiftMessage': parse_gif_msg,  # 礼物消息
                    'WebcastLikeMessage': parse_like_msg,  # 点赞消息
                    'WebcastMemberMessage': parse_member_msg,  # 进入直播间消息
                    'WebcastSocialMessage': parse_social_msg,  # 关注消息
                    # 'WebcastRoomUserSeqMessage': parse_room_user_seq_msg,  # 直播间统计
                    # 'WebcastFansclubMessage': parse_fansclub_msg,  # 粉丝团消息
                    'WebcastControlMessage': self.parse_control_msg,  # 直播间状态消息
                    # 'WebcastEmojiChatMessage': parse_emoji_chat_msg,  # 聊天表情包消息
                    # 'WebcastRoomStatsMessage': parse_room_stats_msg,  # 直播间统计信息
                    # 'WebcastRoomMessage': parse_room_msg,  # 直播间信息
                    # 'WebcastRoomRankMessage': parse_rank_msg,  # 直播间排行榜信息
                    # 'WebcastRoomStreamAdaptationMessage': parse_room_stream_adaptation_msg,  # 直播间流配置
                }
                if method in funcs:
                    funcs[method](msg.payload, self.queue)
            except Exception:
                logger.error(traceback.format_exc())

    def parse_control_msg(self, payload):
        '''直播间状态消息'''
        message = ControlMessage().parse(payload)

        if message.status == 3:
            logger.info("直播间已结束")
            self.stop()

    def _wsOnError(self, ws, error):
        logger.error("WebSocket error: ", error)

    def _wsOnClose(self, ws, *args):
        self.get_room_status()
        logger.info("WebSocket connection closed.")


class DouyinLiveWebReply:
    def __init__(self, queue: PromptQueue, config):
        self.queue = queue
        self.system_text_list = []
        self.backend_token = ''
        self.live_chat_config = LiveChatConfig()
        self.punctuation = ",.!;:，。！？：；"
        self.system_message_index = 0
        self.response_queue = PromptQueue(10)
        self.config = config

    def _llm(self, prompt, stream=False):
        payload = {
            "model": "qwen3:30b-a3b",
            "messages": [
                {
                    "role": "user",
                    "content": f"{prompt}/no_think"
                }
            ],
            "options": {
                "temperature": 0.9
            },
            "stream": False,
            "filterThink": True
        }

        def _gen():
            response = requests.post(
                f'{self.live_chat_config.ollama_address}/live-digital-avatar-manage/ollama/generate', json=payload,
                headers={'Authorization': f'Bearer {self.live_chat_config.backend_token}'},
                stream=True)
            buffer = ''
            for line in response.iter_lines():
                if not line:
                    continue
                logger.info(f'llm output -> {line.decode()}')
                data = json.loads(line.decode()[5:])
                content = data['message']['content']
                for char in content:
                    buffer += char
                    if char in self.punctuation:
                        if len(buffer.strip()) < 10:
                            continue
                        yield buffer.strip()
                        buffer = ''
            if buffer.strip():
                yield buffer.strip()

        if stream:
            return _gen()
        else:
            response = requests.post(
                f'{self.live_chat_config.ollama_address}/live-digital-avatar-manage/ollama/generate', json=payload,
                headers={'Authorization': f'Bearer {self.live_chat_config.backend_token}'},
                timeout=30).content.decode()[5:]
            response = json.loads(response)
            return response['message']['content']

    def reply_message_postprocess(self, text):
        prohibited_words = self.live_chat_config.prohibited_words
        for prohibited_word, substitutes in prohibited_words.items():
            if prohibited_word in text:
                text = text.replace(prohibited_word, substitutes if substitutes is not None else '')
        return text

    async def post_to_human(self, text: str):
        async with httpx.AsyncClient() as client:
            await client.post(
                f'{self.live_chat_config.livetalking_address}/human',
                json={
                    "type": "echo",
                    "sessionid": self.live_chat_config.livetalking_sessionid,
                    "text": text
                },
                timeout=5
            )

    def post_to_human_sync(self, text: str):
        response = requests.post(
            f'{self.live_chat_config.livetalking_address}/human',
            json={
                "type": "echo",
                "sessionid": self.live_chat_config.livetalking_sessionid,
                "text": text
            },
            timeout=5
        )
        if response.status_code != 200:
            logger.error(f'Failed to post to human: {response.text}')

    def generate_messages(self, batch_number=0):
        message_count = self.live_chat_config.messages(batch_number)
        if message_count == 0:
            logger.info(f'生成系统文案，batch_number: {batch_number}')
            system_messages = self.live_chat_config.system_messages
            llm_prompt = self.live_chat_config.refine_system_message_prompt.format(
                content=system_messages)
            logger.info(f'llm_prompt: {llm_prompt}')
            reply_messages = self._llm(llm_prompt, False)
            reply_messages = re.findall(r'\{.*?\}', reply_messages, re.DOTALL)[0]
            reply_messages = json.loads(reply_messages)
            for _type, message in reply_messages.items():
                self.live_chat_config.insert_message(message, _type, batch_number)
                logger.info(f'入库文案:{_type} | {message}')


    def text_to_adio_cost_predict(self, text):
        """
        计算文本转语音的耗时
        :param text: 文本
        :return: 耗时 单位秒
        """
        return len(text)/5

    def __call__(self):
        """
        优先从用户交互队列中取提示词，如果没有用户交互的数据，则输出系统提示词
        """
        # 这里修正为使用 self.live_chat_config
        self.live_chat_config.update_chat_enable_status('已启动')
        logger.info(f'livetalking address -> {self.live_chat_config.livetalking_address}')
        logger.info(f'ollama_address -> {self.live_chat_config.ollama_address}')

        # 下次调用post_to_human的时间
        next_to_human_call_time = None

        while True:
            try:
                # next_to_human_call_time大于当前时间，立即调用
                call_immediate = next_to_human_call_time is None or time.time() >= next_to_human_call_time
                if not call_immediate:
                    prompt_data = self.queue.get(False)
                    if prompt_data is not None:
                        product_name, product_specification, product_description = self.live_chat_config.product_info

                        if isinstance(prompt_data, tuple) and len(prompt_data) == 4:
                            message_type, prompt, live_chat, sender_name = prompt_data
                        else:
                            message_type, prompt, live_chat = prompt_data
                            sender_name = None

                        if message_type == MessageType.ENTER_LIVE_ROOM.value:
                            if random.random() >= self.live_chat_config.enter_live_room_prob / 100:
                                continue
                            else:
                                prompt = prompt.format(product_name=product_name,
                                                       product_specification=product_specification,
                                                       product_description=product_description)
                        elif message_type == MessageType.FOLLOW.value:
                            if random.random() >= self.live_chat_config.follow_prob / 100:
                                continue
                            else:
                                prompt = prompt.format(product_name=product_name,
                                                       product_specification=product_specification,
                                                       product_description=product_description)
                        elif message_type == MessageType.LIKE.value:
                            if random.random() >= self.live_chat_config.like_prob / 100:
                                continue
                            else:
                                prompt = prompt.format(product_name=product_name,
                                                       product_specification=product_specification,
                                                       product_description=product_description)
                        elif message_type == MessageType.GIFT.value:
                            if random.random() >= self.live_chat_config.gift_prob / 100:
                                continue
                            else:
                                prompt = prompt.format(product_name=product_name,
                                                       product_specification=product_specification,
                                                       product_description=product_description)
                        elif message_type == MessageType.CHAT.value:
                            if random.random() >= self.live_chat_config.chat_prob / 100:
                                continue
                            else:
                                prompt = prompt.format(product_name=product_name,
                                                       product_specification=product_specification,
                                                       product_description=product_description)

                        if live_chat is not None:
                            logger.info(f'弹幕({sender_name}): {live_chat}')
                            llm_output = self._llm(
                                self.live_chat_config.product_related_prompt.format(content=live_chat))
                            logger.info(f'判断弹幕是否违反中国大陆法律和政策: {llm_output}')
                            if llm_output != '否':
                                continue

                        reply_message = self._llm(prompt, False)
                        # 这里把“回复内容 + 发送者昵称”一起塞进 response_queue
                        self.response_queue.put((reply_message, sender_name))
                    else:
                        if not self.live_chat_config.next_reply_message:
                            logger.info('文案已用完，重新生成文案')
                            self.live_chat_config.flush_message()
                            self.generate_messages(1)
                    continue
                else:
                    # 不在说话 -> 从队列/库里取一句话播报
                    if not self.response_queue.empty():
                        reply_message, from_user = self.response_queue.get()
                        reply_message = self.reply_message_postprocess(reply_message)
                        if from_user:
                            reply_message = f"{from_user}哥哥，{reply_message}"
                            logger.info(f"开始播放回复弹幕（{from_user}）：{reply_message}")
                        else:
                            logger.info(f"开始播放回复弹幕：{reply_message}")
                    else:
                        precedence_message = self.live_chat_config.precedence_reply_message
                        if not precedence_message:
                            message = self.live_chat_config.next_reply_message
                            if not message:
                                logger.info('文案已用完，重新生成文案')
                                self.generate_messages(0)
                                self.generate_messages(1)
                                continue
                            reply_message, _id = message
                            logger.info(f'更新文案id:{_id}状态为: 1')
                            self.live_chat_config.update_next_reply_status(1, _id)
                        else:
                            reply_message, _id = precedence_message
                            self.live_chat_config.flush_precedence_reply_message()
                        logger.info(f"开始播放系统文案：{reply_message}")
                    cost = self.text_to_adio_cost_predict(reply_message)
                    logger.info(f'文本:{reply_message}预计算出的语音时长：{cost}秒')
                    next_to_human_call_time = time.time() + cost
                    formatted_time = datetime.fromtimestamp(next_to_human_call_time).strftime("%Y-%m-%d %H:%M:%S")
                    logger.info(f'预计在{formatted_time}发送下一段语音')
                    self.post_to_human_sync(reply_message)

            except Exception:
                logger.error(traceback.format_exc())
                time.sleep(5)
                system_messages = self.live_chat_config.system_messages
                reply_message = system_messages[self.system_message_index]
                self.system_message_index += 1
                if self.system_message_index >= len(system_messages):
                    self.system_message_index = 0
                llm_prompt = self.live_chat_config.refine_system_message_prompt.format(content=reply_message)
                reply_messages = self._llm(llm_prompt, True)
                for reply_message in reply_messages:
                    reply_message = self.reply_message_postprocess(reply_message)
                    asyncio.run(self.post_to_human(reply_message))