WeChatMsg/wxManager/manager_v4.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
@Time        : 2024/12/11 20:43 
@Author      : SiYuan 
@Email       : 863909694@qq.com 
@File        : MemoTrace-manager_v4.py
@Description : 
"""
import concurrent
import os
from concurrent.futures import ProcessPoolExecutor, as_completed, ThreadPoolExecutor
from datetime import date, datetime
from multiprocessing import Pool, cpu_count
from typing import Tuple, List, Any

import zstandard as zstd

from wxManager import MessageType
from wxManager.db_v4.biz_message import BizMessageDB
from wxManager.db_v4.emotion import EmotionDB
from wxManager.db_v4.media import MediaDB
from wxManager.db_v4 import ContactDB, HeadImageDB, SessionDB, MessageDB, HardLinkDB
from wxManager.db_main import DataBaseInterface, Context
from wxManager.model.contact import Contact, ContactType, Person
from wxManager.model import Me
from wxManager.parser.util.protocbuf.roomdata_pb2 import ChatRoomData
from wxManager.parser.wechat_v4 import FACTORY_REGISTRY, Singleton
from wxManager.log import logger
from wxManager.parser.util.protocbuf import contact_pb2
from google.protobuf.json_format import MessageToDict


def decompress(data):
    dctx = zstd.ZstdDecompressor()  # 创建解压对象
    x = dctx.decompress(data)
    return x.decode('utf-8')


def parser_messages(messages, username, db_dir=''):
    context = DataBaseV4()
    context.init_database(db_dir)
    if username.endswith('@chatroom'):
        contacts = context.get_chatroom_members(username)
    else:
        contacts = {
            Me().wxid: context.get_contact_by_username(Me().wxid),
            username: context.get_contact_by_username(username)
        }
    # FACTORY_REGISTRY[-1].set_contacts(contacts) # 不知道为什么用对象修改类属性每个实例对象的contacts不一样
    Singleton.set_contacts(contacts)

    for message in messages:
        type_ = message[2]
        if type_ not in FACTORY_REGISTRY:
            type_ = -1
        yield FACTORY_REGISTRY[type_].create(message, username, context)


def _process_messages_batch(messages_batch, username, db_dir) -> List:
    """Helper function to process a batch of messages."""
    processed = []
    for message in parser_messages(messages_batch, username, db_dir):
        processed.append(message)
    return processed


class DataBaseV4(DataBaseInterface):
    def __init__(self):
        super().__init__()
        self.db_dir = ''
        self.chatroom_members_map = {}
        self.contacts_map = {}

        # V4
        self.contact_db = ContactDB('contact/contact.db')
        self.head_image_db = HeadImageDB('head_image/head_image.db')
        self.session_db = SessionDB('session/session.db')
        self.message_db = MessageDB('message/message_0.db', is_series=True)
        self.biz_message_db = BizMessageDB('message/biz_message_0.db', is_series=True)
        self.media_db = MediaDB('message/media_0.db', is_series=True)
        self.hardlink_db = HardLinkDB('hardlink/hardlink.db')
        self.emotion_db = EmotionDB('emoticon/emoticon.db')

    def init_database(self, db_dir=''):
        Me().load_from_json(os.path.join(db_dir, 'info.json'))  # 加载自己的信息
        # print('初始化数据库', db_dir)
        self.db_dir = db_dir
        flag = True
        flag &= self.contact_db.init_database(db_dir)
        flag &= self.head_image_db.init_database(db_dir)
        flag &= self.session_db.init_database(db_dir)
        flag &= self.message_db.init_database(db_dir)
        flag &= self.biz_message_db.init_database(db_dir)
        flag &= self.media_db.init_database(db_dir)
        flag &= self.hardlink_db.init_database(db_dir)
        flag &= self.emotion_db.init_database(db_dir)
        return flag

    def close(self):
        pass

        # self.head_image_db.close()
        # self.contact_db.close()

    def get_session(self):
        """
        获取聊天会话窗口，在聊天界面显示
        @return:
        """
        return self.session_db.get_session()

    def get_messages(
            self,
            username_: str,
            time_range: Tuple[int | float | str | date, int | float | str | date] = None,
    ):
        # todo 改成yield进行操作，多进程处理加快速度
        import time
        st = time.time()
        logger.error(f'开始获取聊天记录：{st}')
        res = []

        # messages = self.message_db.get_messages_by_username(username_, time_range)*20
        # # for messages in self.message_db.get_messages_by_username(username_, time_range):
        # for messages_ in messages:
        #     for message in parser_messages(messages_, username_, self.db_dir):
        #         res.append(message)

        def split_list(lst, n):
            k, m = divmod(len(lst), n)
            return [lst[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)]

        #
        # # # Step 1: Retrieve raw message batches
        if username_.startswith('gh_'):
            messages = self.biz_message_db.get_messages_by_username(username_, time_range)
        else:
            messages = self.message_db.get_messages_by_username(username_, time_range)

        if len(messages) < 20000:
            for message in parser_messages(messages, username_, self.db_dir):
                res.append(message)
        else:
            raw_message_batches = split_list(messages, len(messages) // 10000 + 1)
            #
            # # Step 2: Use multiprocessing to process the message batches
            # res = []
            # for batch in raw_message_batches:
            #     print(len(batch))

            with ProcessPoolExecutor(max_workers=min(len(raw_message_batches), 16)) as executor:
                # Submit tasks
                future_to_batch = {
                    executor.submit(_process_messages_batch, batch, username_, self.db_dir): batch
                    for batch in raw_message_batches
                }

                # Collect results
                for future in future_to_batch.keys():
                    res.extend(future.result())

        et = time.time()
        logger.error(f'获取聊天记录完成：{et}')
        logger.error(f'获取聊天记录耗时：{et - st:.2f}s/{len(res)}条消息 {username_}')
        res.sort()
        return res

    def get_messages_by_num(self, username, start_sort_seq, msg_num=20):
        """
        获取小于start_sort_seq的msg_num个消息
        @param username:
        @param start_sort_seq:
        @param msg_num:
        @return: messages, 最后一条消息的start_sort_seq
        """
        result = []
        if username.startswith('gh_'):
            messages = self.biz_message_db.get_messages_by_num(username, start_sort_seq, msg_num)
        else:
            messages = self.message_db.get_messages_by_num(username, start_sort_seq, msg_num)
        for messages in messages:
            for message in parser_messages(messages, username, self.db_dir):
                result.append(message)
        result.sort(reverse=True)
        res = result[:msg_num]
        return res, res[-1].sort_seq if res else 0

    def get_message_by_server_id(self, username, server_id):
        """
        获取小于start_sort_seq的msg_num个消息
        @param username:
        @param server_id:
        @return: messages, 最后一条消息的start_sort_seq
        """
        message = self.message_db.get_message_by_server_id(username, server_id)
        if message:
            messages_iter = parser_messages([message], username, self.db_dir)
            return next(messages_iter)
        return None

    def get_messages_by_type(
            self,
            username_,
            type_: MessageType,
            time_range: Tuple[int | float | str | date, int | float | str | date] = None,
    ):
        def split_list(lst, n):
            k, m = divmod(len(lst), n)
            return [lst[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)]

        res = []
        # # # Step 1: Retrieve raw message batches
        if username_.startswith('gh_'):
            messages = self.biz_message_db.get_messages_by_type(username_, time_range)
        else:
            messages = self.message_db.get_messages_by_type(username_, type_, time_range)

        if len(messages) < 20000:
            for message in parser_messages(messages, username_, self.db_dir):
                res.append(message)
        else:
            raw_message_batches = split_list(messages, len(messages) // 10000 + 1)
            with ProcessPoolExecutor(max_workers=min(len(raw_message_batches), 16)) as executor:
                # Submit tasks
                future_to_batch = {
                    executor.submit(_process_messages_batch, batch, username_, self.db_dir): batch
                    for batch in raw_message_batches
                }
                # Collect results
                for future in future_to_batch.keys():
                    res.extend(future.result())
        res.sort()
        return res

    def get_messages_calendar(self, username_: str):
        if username_.startswith('gh_'):
            return self.biz_message_db.get_messages_calendar(username_)
        else:
            return self.message_db.get_messages_calendar(username_)

    def get_chatted_top_contacts(
            self,
            time_range: Tuple[int | float | str | date, int | float | str | date] = None,
            contain_chatroom=False,
            top_n=10
    ) -> list:
        return []

    def get_emoji_url(self, md5: str, thumb: bool = False) -> str | bytes:
        return self.emotion_db.get_emoji_url(md5, thumb)

    # 图片、视频、文件
    def get_file(self, md5: bytes | str) -> str:
        return self.hardlink_db.get_file(md5)

    def get_image(self, content, bytesExtra, up_dir="", md5=None, thumb=False, talker_username='') -> str:
        return self.hardlink_db.get_image(content, bytesExtra, up_dir, md5, thumb, talker_username)

    def get_video(self, content, bytesExtra, md5=None, thumb=False):
        return self.hardlink_db.get_video(md5, thumb)

    # 语音
    def get_audio(self, reserved0, output_path, open_im=False, filename=''):
        return self.media_db.get_audio(reserved0, output_path, filename)

    def get_media_buffer(self, server_id, is_open_im=False) -> bytes:
        return self.media_db.get_media_buffer(server_id)

    def get_audio_path(self, reserved0, output_path, filename=''):
        return self.media_db.get_audio_path(reserved0, output_path, filename)

    def get_audio_text(self, msgSvrId):
        return ''

    def update_audio_to_text(self):
        # todo
        return

    # 语音结束

    # 联系人

    def get_avatar_buffer(self, username) -> bytes:
        return self.head_image_db.get_avatar_buffer(username)

    def create_contact(self, contact_info_list) -> Person:
        wxid, local_type, flag = contact_info_list[0], contact_info_list[2], contact_info_list[3]
        nickname = contact_info_list[5]
        remark = contact_info_list[4]
        if not nickname and wxid.endswith('@chatroom'):
            nickname = self._get_chatroom_name(contact_info_list[0])
        if not remark:
            remark = nickname
        gender = '未知'
        signature = ''
        label_list = []
        region = ('', '', '')
        if not (wxid.endswith('@openim') or wxid.endswith('@chatroom')):
            try:
                # 创建顶级消息对象
                message = contact_pb2.ContactInfo()
                # 解析二进制数据
                message.ParseFromString(contact_info_list[10])
                # 转换为 JSON 格式
                detail = MessageToDict(message)
                gender_code = detail.get('gender', 0)
                if gender_code == 1:
                    gender = '男'
                elif gender_code == 2:
                    gender = '女'
                label_list = detail.get('labelList', '').strip(',').split(',')
                signature = detail.get('signature', '')
                region = (detail.get('country', ''), detail.get('province', ''), detail.get('city', ''))
                label_list = self.contact_db.get_labels(detail.get('labelList')).split(',')
            except:
                pass
                # logger.error(f'{wxid} {contact_info_list[5]}联系人解析失败\n{contact_info_list[10]}')
        contact = Contact(
            wxid=contact_info_list[0],
            remark=remark,
            alias=contact_info_list[1],
            nickname=nickname,
            small_head_img_url=contact_info_list[8],
            big_head_img_url=contact_info_list[9],
            flag=contact_info_list[3],
            gender=gender,
            signature=signature,
            label_list=label_list,
            region=region
        )

        def is_nth_bit_set(number, n):
            # 左移 1 到第 n 位
            mask = 1 << n
            # 使用位与运算判断第 n 位
            return (number & mask) != 0

        if local_type == 1:
            contact.type = ContactType.Normal
            if wxid.startswith('gh_'):
                contact.type |= ContactType.Public
            elif wxid.endswith('@chatroom'):
                contact.type |= ContactType.Chatroom
        elif local_type == 2:
            contact.type = ContactType.Chatroom
        elif local_type == 3:
            contact.type = ContactType.Stranger
        elif local_type == 5:
            contact.type = ContactType.OpenIM
        if is_nth_bit_set(flag, 6):
            contact.type |= ContactType.Star
        if is_nth_bit_set(flag, 11):
            contact.type |= ContactType.Sticky

        if local_type == 10086:
            contact.type = ContactType.Unknown
            contact.is_unknown = True
        return contact

    def get_contacts(self) -> List[Person]:
        contacts = []
        contact_lists = self.contact_db.get_contacts()
        for contact_info_list in contact_lists:
            if contact_info_list:
                contact = self.create_contact(contact_info_list)
                contacts.append(contact)
        return contacts

    def set_remark(self, username: str, remark) -> bool:
        if username in self.contacts_map:
            self.contacts_map[username].remark = remark
        return self.contact_db.set_remark(username, remark)

    def set_avatar_buffer(self, username, avatar_path):
        return self.head_image_db.set_avatar_buffer(username, avatar_path)

    def get_contact_by_username(self, wxid: str) -> Person:
        contact_info_list = self.contact_db.get_contact_by_username(wxid)
        if contact_info_list:
            contact = self.create_contact(contact_info_list)
            return contact
        else:
            contact = Contact(
                wxid=wxid,
                nickname=wxid,
                remark=wxid
            )
        return contact

    def get_chatroom_members(self, chatroom_name) -> dict[Any, Person] | Any:
        """
        获取群成员
        @param chatroom_name:
        @return:
        """
        if chatroom_name in self.chatroom_members_map:
            return self.chatroom_members_map[chatroom_name]
        result = {}
        chatroom = self.contact_db.get_chatroom_info(chatroom_name)

        if chatroom is None:
            return result
        # 解析RoomData数据
        parsechatroom = ChatRoomData()
        parsechatroom.ParseFromString(chatroom[1])
        # 群成员数据放入字典存储
        for mem in parsechatroom.members:
            contact = self.get_contact_by_username(mem.wxID)
            if contact:
                if mem.displayName:
                    contact.remark = mem.displayName
                result[contact.wxid] = contact
        self.chatroom_members_map[chatroom_name] = result
        return result

    def _get_chatroom_name(self, wxid):
        chatroom = self.contact_db.get_chatroom_info(wxid)

        if chatroom is None:
            return ''
        # 解析RoomData数据
        parsechatroom = ChatRoomData()
        parsechatroom.ParseFromString(chatroom[1])
        chatroom_name = ''
        # 群成员数据放入字典存储
        for mem in parsechatroom.members[:5]:
            if mem.wxID == Me().wxid:
                continue
            if mem.displayName:
                chatroom_name += f'{mem.displayName}、'
            else:
                contact = self.get_contact_by_username(mem.wxID)
                chatroom_name += f'{contact.remark}、'
        return chatroom_name.rstrip('、')

    # 联系人结束

    def add_audio_txt(self, msgSvrId, text):
        return self.audio_to_text.add_text(msgSvrId, text)

    def get_favorite_items(self, time_range):
        return self.favorite_db.get_items(time_range)

    def merge(self, db_dir):
        """
        批量将db_path中的数据合入到数据库中
        @param db_path:
        @return:
        """
        merge_tasks = {
            self.head_image_db: os.path.join(db_dir, 'head_image', 'head_image.db'),
            self.hardlink_db: os.path.join(db_dir, 'hardlink', 'hardlink.db'),
            self.media_db: os.path.join(db_dir, 'message', 'media_0.db'),
            self.contact_db: os.path.join(db_dir, 'contact', 'contact.db'),
            self.emotion_db: os.path.join(db_dir, 'emoticon', 'emoticon.db'),
            self.message_db: os.path.join(db_dir, 'message', 'message_0.db'),
            self.biz_message_db: os.path.join(db_dir, 'message', 'biz_message_0.db'),
            self.session_db: os.path.join(db_dir, 'session', 'session.db'),
        }

        def merge_task(db_instance, db_path):
            """执行单个数据库的合并任务"""
            db_instance.merge(db_path)

        # 使用 ThreadPoolExecutor 进行多线程合并
        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = {executor.submit(merge_task, db, path): (db, path) for db, path in merge_tasks.items()}

            # 等待所有任务完成
            for future in concurrent.futures.as_completed(futures):
                db, path = futures[future]
                try:
                    future.result()  # 这里会抛出异常（如果有的话）
                    print(f"成功合并数据库: {path}")
                except Exception as e:
                    print(f"合并 {path} 失败: {e}")