LangBot/src/langbot/libs/wecom_ai_bot_api/api.py

import asyncio
import base64
import json
import time
import traceback
import uuid
import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
import re
from typing import Any, Callable, Optional, Tuple
from urllib.parse import unquote

import httpx
from Crypto.Cipher import AES
from quart import Quart, request, Response, jsonify

from langbot.libs.wecom_ai_bot_api import wecombotevent
from langbot.libs.wecom_ai_bot_api.WXBizMsgCrypt3 import WXBizMsgCrypt
from langbot.pkg.platform.logger import EventLogger


@dataclass
class StreamChunk:
    """描述单次推送给企业微信的流式片段。"""

    # 需要返回给企业微信的文本内容
    content: str

    # 标记是否为最终片段，对应企业微信协议里的 finish 字段
    is_final: bool = False

    # 预留额外元信息，未来支持多模态扩展时可使用
    meta: dict[str, Any] = field(default_factory=dict)


@dataclass
class StreamSession:
    """维护一次企业微信流式会话的上下文。"""

    # 企业微信要求的 stream_id，用于标识后续刷新请求
    stream_id: str

    # 原始消息的 msgid，便于与流水线消息对应
    msg_id: str

    # 群聊会话标识（单聊时为空）
    chat_id: Optional[str]

    # 触发消息的发送者
    user_id: Optional[str]

    # 会话创建时间
    created_at: float = field(default_factory=time.time)

    # 最近一次被访问的时间，cleanup 依据该值判断过期
    last_access: float = field(default_factory=time.time)

    # 将流水线增量结果缓存到队列，刷新请求逐条消费
    queue: asyncio.Queue = field(default_factory=asyncio.Queue)

    # 是否已经完成（收到最终片段）
    finished: bool = False

    # 缓存最近一次片段，处理重试或超时兜底
    last_chunk: Optional[StreamChunk] = None

    # 反馈 ID，用于接收用户点赞/点踩反馈
    feedback_id: Optional[str] = None


class StreamSessionManager:
    """管理 stream 会话的生命周期，并负责队列的生产消费。"""

    # Sessions with registered feedback_ids use a longer TTL to survive the
    # full like → cancel → dislike feedback flow. Must align with the adapter's
    # _stream_to_monitoring_msg TTL (wecombot.py).
    _FEEDBACK_SESSION_TTL = 600  # 10 minutes

    def __init__(self, logger: EventLogger, ttl: int = 60) -> None:
        self.logger = logger

        self.ttl = ttl  # 超时时间（秒），超过该时间未被访问的会话会被清理由 cleanup
        self._sessions: dict[str, StreamSession] = {}  # stream_id -> StreamSession 映射
        self._msg_index: dict[str, str] = {}  # msgid -> stream_id 映射，便于流水线根据消息 ID 找到会话
        self._feedback_index: dict[str, str] = {}  # feedback_id -> stream_id 映射

    def get_stream_id_by_msg(self, msg_id: str) -> Optional[str]:
        if not msg_id:
            return None
        return self._msg_index.get(msg_id)

    def get_session(self, stream_id: str) -> Optional[StreamSession]:
        return self._sessions.get(stream_id)

    def get_session_by_feedback_id(self, feedback_id: str) -> Optional[StreamSession]:
        """根据 feedback_id 查找会话。

        Args:
            feedback_id: 企业微信反馈事件中的反馈 ID。

        Returns:
            Optional[StreamSession]: 找到的会话实例，未找到返回 None。
        """
        if not feedback_id:
            return None
        stream_id = self._feedback_index.get(feedback_id)
        if stream_id:
            return self._sessions.get(stream_id)
        return None

    def register_feedback_id(self, stream_id: str, feedback_id: str) -> None:
        """注册 feedback_id 与 stream_id 的映射。

        Args:
            stream_id: 企业微信流式会话 ID。
            feedback_id: 反馈 ID。
        """
        if feedback_id and stream_id:
            self._feedback_index[feedback_id] = stream_id

    def create_or_get(self, msg_json: dict[str, Any]) -> tuple[StreamSession, bool]:
        """根据企业微信回调创建或获取会话。

        Args:
            msg_json: 企业微信解密后的回调 JSON。

        Returns:
            Tuple[StreamSession, bool]: `StreamSession` 为会话实例，`bool` 指示是否为新建会话。

        Example:
            在首次回调中调用，得到 `is_new=True` 后再触发流水线。
        """
        msg_id = msg_json.get('msgid', '')
        if msg_id and msg_id in self._msg_index:
            stream_id = self._msg_index[msg_id]
            session = self._sessions.get(stream_id)
            if session:
                session.last_access = time.time()
                return session, False

        stream_id = str(uuid.uuid4())
        session = StreamSession(
            stream_id=stream_id,
            msg_id=msg_id,
            chat_id=msg_json.get('chatid'),
            user_id=msg_json.get('from', {}).get('userid'),
        )

        if msg_id:
            self._msg_index[msg_id] = stream_id
        self._sessions[stream_id] = session
        return session, True

    async def publish(self, stream_id: str, chunk: StreamChunk) -> bool:
        """向 stream 队列写入新的增量片段。

        Args:
            stream_id: 企业微信分配的流式会话 ID。
            chunk: 待发送的增量片段。

        Returns:
            bool: 当流式队列存在并成功入队时返回 True。

        Example:
            在收到模型增量后调用 `await manager.publish('sid', StreamChunk('hello'))`。
        """
        session = self._sessions.get(stream_id)
        if not session:
            return False

        session.last_access = time.time()
        session.last_chunk = chunk

        try:
            session.queue.put_nowait(chunk)
        except asyncio.QueueFull:
            # 默认无界队列，此处兜底防御
            await session.queue.put(chunk)

        if chunk.is_final:
            session.finished = True

        return True

    async def consume(self, stream_id: str, timeout: float = 0.5) -> Optional[StreamChunk]:
        """从队列中取出一个片段，若超时返回 None。

        Args:
            stream_id: 企业微信流式会话 ID。
            timeout: 取片段的最长等待时间（秒）。

        Returns:
            Optional[StreamChunk]: 成功时返回片段，超时或会话不存在时返回 None。

        Example:
            企业微信刷新到达时调用，若队列有数据则立即返回 `StreamChunk`。
        """
        session = self._sessions.get(stream_id)
        if not session:
            return None

        session.last_access = time.time()

        try:
            chunk = await asyncio.wait_for(session.queue.get(), timeout)
            session.last_access = time.time()
            if chunk.is_final:
                session.finished = True
            return chunk
        except asyncio.TimeoutError:
            if session.finished and session.last_chunk:
                return session.last_chunk
            return None

    def mark_finished(self, stream_id: str) -> None:
        session = self._sessions.get(stream_id)
        if session:
            session.finished = True
            session.last_access = time.time()

    def cleanup(self) -> None:
        """定期清理过期会话，防止队列与映射无上限累积。

        已注册 feedback_id 的会话使用更长的 TTL，确保用户在点赞/取消/点踩流程中
        不会因为 session 被提前清除而丢失上下文信息。
        """
        now = time.time()
        expired: list[str] = []
        for stream_id, session in self._sessions.items():
            # Sessions with registered feedback_ids use a longer TTL
            effective_ttl = self._FEEDBACK_SESSION_TTL if session.feedback_id else self.ttl
            if now - session.last_access > effective_ttl:
                expired.append(stream_id)

        for stream_id in expired:
            session = self._sessions.pop(stream_id, None)
            if not session:
                continue
            msg_id = session.msg_id
            if msg_id and self._msg_index.get(msg_id) == stream_id:
                self._msg_index.pop(msg_id, None)
            # Clean up feedback index for expired sessions
            if session.feedback_id:
                self._feedback_index.pop(session.feedback_id, None)


def _decrypt_file(encrypted_data: bytes, aes_key_str: str) -> bytes:
    """Decrypt AES-256-CBC encrypted file data.

    Aligned with the official WeCom AI Bot Python SDK (crypto_utils.py).

    Args:
        encrypted_data: The raw encrypted bytes.
        aes_key_str: Base64-encoded AES key (may lack padding).

    Returns:
        Decrypted bytes with PKCS#7 padding removed.
    """
    if not encrypted_data:
        raise ValueError('encrypted_data is empty')
    if not aes_key_str:
        raise ValueError('aes_key is empty')

    # Python's base64.b64decode requires proper padding (length % 4 == 0).
    # Node.js Buffer.from tolerates missing '=', so we must pad manually.
    remainder = len(aes_key_str) % 4
    if remainder != 0:
        aes_key_str = aes_key_str + '=' * (4 - remainder)
    key = base64.b64decode(aes_key_str)

    iv = key[:16]

    cipher = AES.new(key, AES.MODE_CBC, iv)

    # Ensure encrypted data is aligned to AES block size (16 bytes).
    # Node.js setAutoPadding(false) silently handles unaligned data,
    # but PyCryptodome will raise an error.
    block_size = 16
    data_remainder = len(encrypted_data) % block_size
    if data_remainder != 0:
        encrypted_data = encrypted_data + b'\x00' * (block_size - data_remainder)

    decrypted = cipher.decrypt(encrypted_data)

    # Remove PKCS#7 padding with validation
    if len(decrypted) == 0:
        raise ValueError('Decrypted data is empty')

    pad_len = decrypted[-1]
    if pad_len < 1 or pad_len > 32 or pad_len > len(decrypted):
        raise ValueError(f'Invalid PKCS#7 padding value: {pad_len}')

    # Verify all padding bytes are consistent
    for i in range(len(decrypted) - pad_len, len(decrypted)):
        if decrypted[i] != pad_len:
            raise ValueError('Invalid PKCS#7 padding: padding bytes mismatch')

    return decrypted[: len(decrypted) - pad_len]


def _extract_filename(content_disposition: str) -> Optional[str]:
    """Extract filename from a Content-Disposition header value."""
    if not content_disposition:
        return None
    # RFC 5987: filename*=UTF-8''xxx
    utf8_match = re.search(r"filename\*=UTF-8''([^;\s]+)", content_disposition, re.IGNORECASE)
    if utf8_match:
        return unquote(utf8_match.group(1))
    # Standard: filename="xxx" or filename=xxx
    match = re.search(r'filename="?([^";\s]+)"?', content_disposition, re.IGNORECASE)
    if match:
        return unquote(match.group(1))
    return None


def _bytes_to_data_uri(data: bytes) -> str:
    """Convert raw bytes to a data URI with auto-detected MIME type."""
    if data.startswith(b'\xff\xd8'):
        mime_type = 'image/jpeg'
    elif data.startswith(b'\x89PNG'):
        mime_type = 'image/png'
    elif data.startswith((b'GIF87a', b'GIF89a')):
        mime_type = 'image/gif'
    elif data.startswith(b'BM'):
        mime_type = 'image/bmp'
    elif data.startswith(b'II*\x00') or data.startswith(b'MM\x00*'):
        mime_type = 'image/tiff'
    elif data[:4] == b'%PDF':
        mime_type = 'application/pdf'
    elif data[:4] == b'PK\x03\x04':
        mime_type = 'application/zip'
    else:
        mime_type = 'application/octet-stream'

    base64_str = base64.b64encode(data).decode('utf-8')
    return f'data:{mime_type};base64,{base64_str}'


async def download_encrypted_file(
    download_url: str, aes_key: str, logger: EventLogger
) -> Tuple[Optional[bytes], Optional[str]]:
    """Download an AES-encrypted file from WeChat Work and decrypt it.

    Args:
        download_url: The encrypted file download URL.
        aes_key: The AES key for decryption (base64-encoded, per-message aeskey
                 or platform EncodingAESKey).
        logger: Logger instance.

    Returns:
        A tuple of (decrypted_bytes, filename) or (None, None) on failure.
    """
    if not download_url:
        return None, None
    if not aes_key:
        await logger.error('download_encrypted_file: aes_key is empty, cannot decrypt')
        return None, None

    filename: Optional[str] = None
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(download_url)
            if response.status_code != 200:
                await logger.error(f'Failed to download file (HTTP {response.status_code}): {response.text[:200]}')
                return None, None
            encrypted_bytes = response.content
            filename = _extract_filename(response.headers.get('content-disposition', ''))
    except Exception:
        await logger.error(f'Failed to download file: {traceback.format_exc()}')
        return None, None

    try:
        decrypted = _decrypt_file(encrypted_bytes, aes_key)
        return decrypted, filename
    except Exception:
        await logger.error(f'Failed to decrypt file: {traceback.format_exc()}')
        return None, None


async def parse_wecom_bot_message(
    msg_json: dict[str, Any], encoding_aes_key: str, logger: EventLogger
) -> dict[str, Any]:
    """Parse a decrypted WeChat Work AI Bot message JSON into a unified message dict.

    This is the shared message parsing logic used by both webhook and WebSocket modes.

    Args:
        msg_json: The decrypted message JSON from WeChat Work.
        encoding_aes_key: AES key for file decryption.
        logger: Logger instance.

    Returns:
        A dict suitable for constructing a WecomBotEvent.
    """
    message_data: dict[str, Any] = {}

    msg_type = msg_json.get('msgtype', '')
    if msg_type:
        message_data['msgtype'] = msg_type

    if msg_json.get('chattype', '') == 'single':
        message_data['type'] = 'single'
    elif msg_json.get('chattype', '') == 'group':
        message_data['type'] = 'group'

    max_inline_file_size = 5 * 1024 * 1024

    async def _safe_download(url: str, per_msg_aeskey: str = '') -> Tuple[Optional[bytes], Optional[str]]:
        """Download and decrypt a file, preferring per-message aeskey over platform key."""
        if not url:
            return None, None
        key = per_msg_aeskey or encoding_aes_key
        if not key:
            await logger.warning('No AES key available for file decryption, skipping download')
            return None, None
        return await download_encrypted_file(url, key, logger)

    async def _safe_download_as_data_uri(url: str, per_msg_aeskey: str = '') -> Optional[str]:
        """Download, decrypt, and convert to data URI for backward compatibility."""
        data, _filename = await _safe_download(url, per_msg_aeskey)
        if data:
            return _bytes_to_data_uri(data)
        return None

    if msg_type == 'text':
        message_data['content'] = msg_json.get('text', {}).get('content')
    elif msg_type == 'markdown':
        message_data['content'] = msg_json.get('markdown', {}).get('content') or msg_json.get('text', {}).get(
            'content', ''
        )
    elif msg_type == 'image':
        image_info = msg_json.get('image', {})
        picurl = image_info.get('url', '')
        per_msg_aeskey = image_info.get('aeskey', '')
        base64_data = await _safe_download_as_data_uri(picurl, per_msg_aeskey)
        if base64_data:
            message_data['picurl'] = base64_data
            message_data['images'] = [base64_data]
    elif msg_type == 'voice':
        voice_info = msg_json.get('voice', {}) or {}
        download_url = voice_info.get('url')
        per_msg_aeskey = voice_info.get('aeskey', '')
        message_data['voice'] = {
            'url': download_url,
            'md5sum': voice_info.get('md5sum') or voice_info.get('md5'),
            'filesize': voice_info.get('filesize') or voice_info.get('size'),
            'sdkfileid': voice_info.get('sdkfileid') or voice_info.get('fileid'),
        }
        if voice_info.get('content'):
            message_data['content'] = voice_info.get('content')
        # if (message_data['voice'].get('filesize') or 0) <= max_inline_file_size:
        #     voice_base64 = await _safe_download_as_data_uri(download_url, per_msg_aeskey)
        #     if voice_base64:
        #         message_data['voice']['base64'] = voice_base64
    elif msg_type == 'video':
        video_info = msg_json.get('video', {}) or {}
        download_url = video_info.get('url')
        per_msg_aeskey = video_info.get('aeskey', '')
        video_data = {
            'url': download_url,
            'filesize': video_info.get('filesize') or video_info.get('size'),
            'sdkfileid': video_info.get('sdkfileid') or video_info.get('fileid'),
            'md5sum': video_info.get('md5sum') or video_info.get('md5'),
            'filename': video_info.get('filename') or video_info.get('name'),
        }
        # if (video_data.get('filesize') or 0) <= max_inline_file_size:
        #     video_base64 = await _safe_download_as_data_uri(download_url, per_msg_aeskey)
        #     if video_base64:
        #         video_data['base64'] = video_base64
        # 应为需要解密，但是目前暂时不能下载到内部进行解密，所以先将下载链接拼接aeskey返回给用户，由插件去处理该链接的下载和解密逻辑
        video_data['download_url'] = download_url + f'?aeskey={per_msg_aeskey}'
        message_data['video'] = video_data
    elif msg_type == 'file':
        file_info = msg_json.get('file', {}) or {}
        download_url = file_info.get('url') or file_info.get('fileurl')
        per_msg_aeskey = file_info.get('aeskey', '')
        file_data = {
            'filename': file_info.get('filename') or file_info.get('name'),
            'filesize': file_info.get('filesize') or file_info.get('size'),
            'md5sum': file_info.get('md5sum') or file_info.get('md5'),
            'sdkfileid': file_info.get('sdkfileid') or file_info.get('fileid'),
            'download_url': download_url,
            'extra': file_info,
        }
        # if (file_data.get('filesize') or 0) <= max_inline_file_size:
        #     file_bytes, dl_filename = await _safe_download(download_url, per_msg_aeskey)
        #     if file_bytes:
        #         file_data['base64'] = _bytes_to_data_uri(file_bytes)
        #         if dl_filename and not file_data.get('filename'):
        #             file_data['filename'] = dl_filename

        # 应为需要解密，但是目前暂时不能下载到内部进行解密，所以先将下载链接拼接aeskey返回给用户，由插件去处理该链接的下载和解密逻辑
        file_data['download_url'] = download_url + f'?aeskey={per_msg_aeskey}'
        message_data['file'] = file_data
    elif msg_type == 'link':
        message_data['link'] = msg_json.get('link', {})
        if not message_data.get('content'):
            title = message_data['link'].get('title', '')
            desc = message_data['link'].get('description') or message_data['link'].get('digest', '')
            message_data['content'] = '\n'.join(filter(None, [title, desc]))
    elif msg_type == 'mixed':
        items = msg_json.get('mixed', {}).get('msg_item', [])
        texts = []
        images = []
        files = []
        voices = []
        videos = []
        links = []
        for item in items:
            item_type = item.get('msgtype')
            if item_type == 'text':
                texts.append(item.get('text', {}).get('content', ''))
            elif item_type == 'image':
                img_info = item.get('image', {})
                img_url = img_info.get('url')
                img_aeskey = img_info.get('aeskey', '')
                base64_data = await _safe_download_as_data_uri(img_url, img_aeskey)
                if base64_data:
                    images.append(base64_data)
            elif item_type == 'file':
                file_info = item.get('file', {}) or {}
                download_url = file_info.get('url') or file_info.get('fileurl')
                item_aeskey = file_info.get('aeskey', '')
                file_data = {
                    'filename': file_info.get('filename') or file_info.get('name'),
                    'filesize': file_info.get('filesize') or file_info.get('size'),
                    'md5sum': file_info.get('md5sum') or file_info.get('md5'),
                    'sdkfileid': file_info.get('sdkfileid') or file_info.get('fileid'),
                    'download_url': download_url,
                    'extra': file_info,
                }
                if (file_data.get('filesize') or 0) <= max_inline_file_size:
                    file_bytes, dl_filename = await _safe_download(download_url, item_aeskey)
                    if file_bytes:
                        file_data['base64'] = _bytes_to_data_uri(file_bytes)
                        if dl_filename and not file_data.get('filename'):
                            file_data['filename'] = dl_filename
                files.append(file_data)
            elif item_type == 'voice':
                voice_info = item.get('voice', {}) or {}
                download_url = voice_info.get('url')
                item_aeskey = voice_info.get('aeskey', '')
                voice_data = {
                    'url': download_url,
                    'md5sum': voice_info.get('md5sum') or voice_info.get('md5'),
                    'filesize': voice_info.get('filesize') or voice_info.get('size'),
                    'sdkfileid': voice_info.get('sdkfileid') or voice_info.get('fileid'),
                }
                if voice_info.get('content'):
                    texts.append(voice_info.get('content'))
                if (voice_data.get('filesize') or 0) <= max_inline_file_size:
                    voice_base64 = await _safe_download_as_data_uri(download_url, item_aeskey)
                    if voice_base64:
                        voice_data['base64'] = voice_base64
                voices.append(voice_data)
            elif item_type == 'video':
                video_info = item.get('video', {}) or {}
                download_url = video_info.get('url')
                item_aeskey = video_info.get('aeskey', '')
                video_data = {
                    'url': download_url,
                    'filesize': video_info.get('filesize') or video_info.get('size'),
                    'sdkfileid': video_info.get('sdkfileid') or video_info.get('fileid'),
                    'md5sum': video_info.get('md5sum') or video_info.get('md5'),
                    'filename': video_info.get('filename') or video_info.get('name'),
                }
                if (video_data.get('filesize') or 0) <= max_inline_file_size:
                    video_base64 = await _safe_download_as_data_uri(download_url, item_aeskey)
                    if video_base64:
                        video_data['base64'] = video_base64
                videos.append(video_data)
            elif item_type == 'link':
                links.append(item.get('link', {}))

        if texts:
            message_data['content'] = ' '.join(texts)
        if images:
            message_data['images'] = images
            message_data['picurl'] = images[0]
        if files:
            message_data['files'] = files
            message_data['file'] = files[0]
        if voices:
            message_data['voices'] = voices
            message_data['voice'] = voices[0]
        if videos:
            message_data['videos'] = videos
            message_data['video'] = videos[0]
        if links:
            message_data['link'] = links[0]
        if items:
            message_data['attachments'] = items
    else:
        message_data['raw_msg'] = msg_json

    from_info = msg_json.get('from', {})
    message_data['userid'] = from_info.get('userid', '')
    message_data['username'] = from_info.get('alias', '') or from_info.get('name', '') or from_info.get('userid', '')

    if msg_json.get('chattype', '') == 'group':
        message_data['chatid'] = msg_json.get('chatid', '')
        message_data['chatname'] = msg_json.get('chatname', '') or msg_json.get('chatid', '')

    message_data['msgid'] = msg_json.get('msgid', '')

    if msg_json.get('aibotid'):
        message_data['aibotid'] = msg_json.get('aibotid', '')

    # Handle quote (referenced message) - important for group chat file references
    quote_info = msg_json.get('quote')
    if quote_info:
        quote_data: dict[str, Any] = {}
        quote_type = quote_info.get('msgtype', '')
        quote_data['msgtype'] = quote_type

        if quote_type == 'text':
            quote_data['content'] = quote_info.get('text', {}).get('content', '')
        elif quote_type == 'image':
            img_info = quote_info.get('image', {})
            img_url = img_info.get('url', '')
            img_aeskey = img_info.get('aeskey', '')
            base64_data = await _safe_download_as_data_uri(img_url, img_aeskey)
            if base64_data:
                quote_data['picurl'] = base64_data
                quote_data['images'] = [base64_data]
        elif quote_type == 'file':
            file_info = quote_info.get('file', {}) or {}
            download_url = file_info.get('url') or file_info.get('fileurl')
            item_aeskey = file_info.get('aeskey', '')
            file_data = {
                'filename': file_info.get('filename') or file_info.get('name'),
                'filesize': file_info.get('filesize') or file_info.get('size'),
                'md5sum': file_info.get('md5sum') or file_info.get('md5'),
                'sdkfileid': file_info.get('sdkfileid') or file_info.get('fileid'),
                'download_url': download_url,
                'extra': file_info,
            }
            # Same as private chat: append aeskey to download_url for plugin processing
            if download_url and item_aeskey:
                file_data['download_url'] = download_url + f'?aeskey={item_aeskey}'
            quote_data['file'] = file_data
        elif quote_type == 'voice':
            voice_info = quote_info.get('voice', {}) or {}
            download_url = voice_info.get('url')
            item_aeskey = voice_info.get('aeskey', '')
            voice_data = {
                'url': download_url,
                'md5sum': voice_info.get('md5sum') or voice_info.get('md5'),
                'filesize': voice_info.get('filesize') or voice_info.get('size'),
                'sdkfileid': voice_info.get('sdkfileid') or voice_info.get('fileid'),
            }
            if voice_info.get('content'):
                quote_data['content'] = voice_info.get('content')
            # Same as private chat: append aeskey to url for plugin processing
            if download_url and item_aeskey:
                voice_data['url'] = download_url + f'?aeskey={item_aeskey}'
            quote_data['voice'] = voice_data
        elif quote_type == 'video':
            video_info = quote_info.get('video', {}) or {}
            download_url = video_info.get('url')
            item_aeskey = video_info.get('aeskey', '')
            video_data = {
                'url': download_url,
                'filesize': video_info.get('filesize') or video_info.get('size'),
                'sdkfileid': video_info.get('sdkfileid') or video_info.get('fileid'),
                'md5sum': video_info.get('md5sum') or video_info.get('md5'),
                'filename': video_info.get('filename') or video_info.get('name'),
            }
            # Same as private chat: append aeskey to download_url for plugin processing
            if download_url and item_aeskey:
                video_data['download_url'] = download_url + f'?aeskey={item_aeskey}'
            quote_data['video'] = video_data
        elif quote_type == 'link':
            quote_data['link'] = quote_info.get('link', {})
            link = quote_data['link']
            title = link.get('title', '')
            desc = link.get('description') or link.get('digest', '')
            quote_data['content'] = '\n'.join(filter(None, [title, desc]))
        elif quote_type == 'mixed':
            # Handle mixed type in quote (text + images + files etc.)
            items = quote_info.get('mixed', {}).get('msg_item', [])
            texts = []
            images = []
            files = []
            for item in items:
                item_type = item.get('msgtype')
                if item_type == 'text':
                    texts.append(item.get('text', {}).get('content', ''))
                elif item_type == 'image':
                    img_info = item.get('image', {})
                    img_url = img_info.get('url')
                    img_aeskey = img_info.get('aeskey', '')
                    base64_data = await _safe_download_as_data_uri(img_url, img_aeskey)
                    if base64_data:
                        images.append(base64_data)
                elif item_type == 'file':
                    file_info = item.get('file', {}) or {}
                    download_url = file_info.get('url') or file_info.get('fileurl')
                    item_aeskey = file_info.get('aeskey', '')
                    file_data = {
                        'filename': file_info.get('filename') or file_info.get('name'),
                        'filesize': file_info.get('filesize') or file_info.get('size'),
                        'md5sum': file_info.get('md5sum') or file_info.get('md5'),
                        'sdkfileid': file_info.get('sdkfileid') or file_info.get('fileid'),
                        'download_url': download_url,
                        'extra': file_info,
                    }
                    # Same as private chat: append aeskey to download_url for plugin processing
                    if download_url and item_aeskey:
                        file_data['download_url'] = download_url + f'?aeskey={item_aeskey}'
                    files.append(file_data)
            if texts:
                quote_data['content'] = ' '.join(texts)
            if images:
                quote_data['images'] = images
                quote_data['picurl'] = images[0]
            if files:
                quote_data['files'] = files
                quote_data['file'] = files[0]

        message_data['quote'] = quote_data

    return message_data


class WecomBotClient:
    def __init__(self, Token: str, EnCodingAESKey: str, Corpid: str, logger: EventLogger, unified_mode: bool = False):
        """企业微信智能机器人客户端。

        Args:
            Token: 企业微信回调验证使用的 token。
            EnCodingAESKey: 企业微信消息加解密密钥。
            Corpid: 企业 ID。
            logger: 日志记录器。
            unified_mode: 是否使用统一 webhook 模式（默认 False）。

        Example:
            >>> client = WecomBotClient(Token='token', EnCodingAESKey='aeskey', Corpid='corp', logger=logger)
        """

        self.Token = Token
        self.EnCodingAESKey = EnCodingAESKey
        self.Corpid = Corpid
        self.ReceiveId = ''
        self.unified_mode = unified_mode
        self.app = Quart(__name__)

        # 只有在非统一模式下才注册独立路由
        if not self.unified_mode:
            self.app.add_url_rule(
                '/callback/command', 'handle_callback', self.handle_callback_request, methods=['POST', 'GET']
            )

        self._message_handlers = {
            'example': [],
        }
        self.logger = logger
        self.generated_content: dict[str, str] = {}
        self.msg_id_map: dict[str, int] = {}
        self.stream_sessions = StreamSessionManager(logger=logger)
        self.stream_poll_timeout = 0.5

        self._feedback_callback: Optional[Callable] = None

    def set_feedback_callback(self, callback: Callable) -> None:
        """设置反馈回调函数。

        Args:
            callback: 反馈回调函数，签名: async def callback(feedback_id, feedback_type, feedback_content, inaccurate_reasons, session)
        """
        self._feedback_callback = callback

    @staticmethod
    def _build_stream_payload(
        stream_id: str, content: str, finish: bool, feedback_id: Optional[str] = None
    ) -> dict[str, Any]:
        """按照企业微信协议拼装返回报文。

        Args:
            stream_id: 企业微信会话 ID。
            content: 推送的文本内容。
            finish: 是否为最终片段。
            feedback_id: 反馈 ID，用于接收用户点赞/点踩反馈。

        Returns:
            dict[str, Any]: 可直接加密返回的 payload。

        Example:
            组装 `{'msgtype': 'stream', 'stream': {'id': 'sid', ...}}` 结构。
        """
        stream_payload = {
            'id': stream_id,
            'finish': finish,
            'content': content,
        }
        if feedback_id:
            stream_payload['feedback'] = {'id': feedback_id}
        return {
            'msgtype': 'stream',
            'stream': stream_payload,
        }

    async def _encrypt_and_reply(self, payload: dict[str, Any], nonce: str) -> tuple[Response, int]:
        """对响应进行加密封装并返回给企业微信。

        Args:
            payload: 待加密的响应内容。
            nonce: 企业微信回调参数中的 nonce。

        Returns:
            Tuple[Response, int]: Quart Response 对象及状态码。

        Example:
            在首包或刷新场景中调用以生成加密响应。
        """
        reply_plain_str = json.dumps(payload, ensure_ascii=False)
        reply_timestamp = str(int(time.time()))
        ret, encrypt_text = self.wxcpt.EncryptMsg(reply_plain_str, nonce, reply_timestamp)
        if ret != 0:
            await self.logger.error(f'加密失败: {ret}')
            return jsonify({'error': 'encrypt_failed'}), 500

        root = ET.fromstring(encrypt_text)
        encrypt = root.find('Encrypt').text
        resp = {
            'encrypt': encrypt,
        }
        return jsonify(resp), 200

    async def _dispatch_event(self, event: wecombotevent.WecomBotEvent) -> None:
        """异步触发流水线处理，避免阻塞首包响应。

        Args:
            event: 由企业微信消息转换的内部事件对象。
        """
        try:
            await self._handle_message(event)
        except Exception:
            await self.logger.error(traceback.format_exc())

    async def _handle_post_initial_response(self, msg_json: dict[str, Any], nonce: str) -> tuple[Response, int]:
        """处理企业微信首次推送的消息，返回 stream_id 并开启流水线。

        Args:
            msg_json: 解密后的企业微信消息 JSON。
            nonce: 企业微信回调参数 nonce。

        Returns:
            Tuple[Response, int]: Quart Response 及状态码。

        Example:
            首次回调时调用，立即返回带 `stream_id` 的响应。
        """
        session, is_new = self.stream_sessions.create_or_get(msg_json)

        feedback_id = str(uuid.uuid4())
        session.feedback_id = feedback_id
        self.stream_sessions.register_feedback_id(session.stream_id, feedback_id)

        message_data = await self.get_message(msg_json)
        if message_data:
            message_data['stream_id'] = session.stream_id
            message_data['feedback_id'] = feedback_id
            try:
                event = wecombotevent.WecomBotEvent(message_data)
            except Exception:
                await self.logger.error(traceback.format_exc())
            else:
                if is_new:
                    asyncio.create_task(self._dispatch_event(event))

        payload = self._build_stream_payload(session.stream_id, '', False, feedback_id)
        return await self._encrypt_and_reply(payload, nonce)

    async def _handle_post_followup_response(self, msg_json: dict[str, Any], nonce: str) -> tuple[Response, int]:
        """处理企业微信的流式刷新请求，按需返回增量片段。

        Args:
            msg_json: 解密后的企业微信刷新请求。
            nonce: 企业微信回调参数 nonce。

        Returns:
            Tuple[Response, int]: Quart Response 及状态码。

        Example:
            在刷新请求中调用，按需返回增量片段。
        """
        stream_info = msg_json.get('stream', {})
        stream_id = stream_info.get('id', '')
        if not stream_id:
            await self.logger.error('刷新请求缺少 stream.id')
            return await self._encrypt_and_reply(self._build_stream_payload('', '', True), nonce)

        session = self.stream_sessions.get_session(stream_id)
        chunk = await self.stream_sessions.consume(stream_id, timeout=self.stream_poll_timeout)

        if not chunk:
            cached_content = None
            if session and session.msg_id:
                cached_content = self.generated_content.pop(session.msg_id, None)
            if cached_content is not None:
                chunk = StreamChunk(content=cached_content, is_final=True)
            else:
                payload = self._build_stream_payload(stream_id, '', False)
                return await self._encrypt_and_reply(payload, nonce)

        payload = self._build_stream_payload(stream_id, chunk.content, chunk.is_final)
        if chunk.is_final:
            self.stream_sessions.mark_finished(stream_id)
        return await self._encrypt_and_reply(payload, nonce)

    async def handle_callback_request(self):
        """企业微信回调入口（独立端口模式，使用全局 request）。

        Returns:
            Quart Response: 根据请求类型返回验证、首包或刷新结果。

        Example:
            作为 Quart 路由处理函数直接注册并使用。
        """
        return await self._handle_callback_internal(request)

    async def handle_unified_webhook(self, req):
        """处理回调请求（统一 webhook 模式，显式传递 request）。

        Args:
            req: Quart Request 对象

        Returns:
            响应数据
        """
        return await self._handle_callback_internal(req)

    async def _handle_callback_internal(self, req):
        """处理回调请求的内部实现，包括 GET 验证和 POST 消息接收。

        Args:
            req: Quart Request 对象
        """
        try:
            self.wxcpt = WXBizMsgCrypt(self.Token, self.EnCodingAESKey, '')

            if req.method == 'GET':
                return await self._handle_get_callback(req)

            if req.method == 'POST':
                return await self._handle_post_callback(req)

            return Response('', status=405)

        except Exception:
            await self.logger.error(traceback.format_exc())
            return Response('Internal Server Error', status=500)

    async def _handle_get_callback(self, req) -> tuple[Response, int] | Response:
        """处理企业微信的 GET 验证请求。"""

        msg_signature = unquote(req.args.get('msg_signature', ''))
        timestamp = unquote(req.args.get('timestamp', ''))
        nonce = unquote(req.args.get('nonce', ''))
        echostr = unquote(req.args.get('echostr', ''))

        if not all([msg_signature, timestamp, nonce, echostr]):
            await self.logger.error('请求参数缺失')
            return Response('缺少参数', status=400)

        ret, decrypted_str = self.wxcpt.VerifyURL(msg_signature, timestamp, nonce, echostr)
        if ret != 0:
            await self.logger.error('验证URL失败')
            return Response('验证失败', status=403)

        return Response(decrypted_str, mimetype='text/plain')

    async def _handle_post_callback(self, req) -> tuple[Response, int] | Response:
        """处理企业微信的 POST 回调请求。"""

        self.stream_sessions.cleanup()

        msg_signature = unquote(req.args.get('msg_signature', ''))
        timestamp = unquote(req.args.get('timestamp', ''))
        nonce = unquote(req.args.get('nonce', ''))

        encrypted_json = await req.get_json()
        encrypted_msg = (encrypted_json or {}).get('encrypt', '')
        if not encrypted_msg:
            await self.logger.error("请求体中缺少 'encrypt' 字段")
            return Response('Bad Request', status=400)

        xml_post_data = f'<xml><Encrypt><![CDATA[{encrypted_msg}]]></Encrypt></xml>'
        ret, decrypted_xml = self.wxcpt.DecryptMsg(xml_post_data, msg_signature, timestamp, nonce)
        if ret != 0:
            await self.logger.error('解密失败')
            return Response('解密失败', status=400)

        msg_json = json.loads(decrypted_xml)

        event = msg_json.get('event', {})
        event_type = event.get('eventtype', '')

        if event_type == 'feedback_event':
            return await self._handle_feedback_event(msg_json, nonce)

        if msg_json.get('msgtype') == 'stream':
            return await self._handle_post_followup_response(msg_json, nonce)

        return await self._handle_post_initial_response(msg_json, nonce)

    async def _handle_feedback_event(self, msg_json: dict[str, Any], nonce: str) -> tuple[Response, int]:
        """处理企业微信用户反馈事件（点赞/点踩）。

        Args:
            msg_json: 解密后的企业微信反馈事件 JSON。
            nonce: 企业微信回调参数 nonce。

        Returns:
            Tuple[Response, int]: Quart Response 及状态码。

        Note:
            企业微信协议要求：反馈事件目前仅支持回复空包。
        """
        try:
            feedback_event = msg_json.get('event', {}).get('feedback_event', {})
            feedback_id = feedback_event.get('id', '')
            feedback_type = feedback_event.get('type', 0)
            feedback_content = feedback_event.get('content', '')
            inaccurate_reasons = feedback_event.get('inaccurate_reason_list', [])

            await self.logger.info(
                f'收到用户反馈事件: feedback_id={feedback_id}, type={feedback_type}, '
                f'content={feedback_content}, reasons={inaccurate_reasons}'
            )

            session = self.stream_sessions.get_session_by_feedback_id(feedback_id)

            if session:
                await self.logger.info(
                    f'反馈关联到会话: stream_id={session.stream_id}, msg_id={session.msg_id}, user_id={session.user_id}'
                )
            else:
                await self.logger.warning(f'未找到 feedback_id={feedback_id} 对应的会话，仍将记录反馈')

            # Dispatch feedback event regardless of session availability
            for handler in self._message_handlers.get('feedback', []):
                try:
                    await handler(
                        feedback_id=feedback_id,
                        feedback_type=feedback_type,
                        feedback_content=feedback_content,
                        inaccurate_reasons=inaccurate_reasons,
                        session=session,
                    )
                except Exception:
                    await self.logger.error(traceback.format_exc())

            if self._feedback_callback:
                try:
                    await self._feedback_callback(
                        feedback_id=feedback_id,
                        feedback_type=feedback_type,
                        feedback_content=feedback_content,
                        inaccurate_reasons=inaccurate_reasons,
                        session=session,
                    )
                except Exception:
                    await self.logger.error(traceback.format_exc())

        except Exception:
            await self.logger.error(traceback.format_exc())

        return await self._encrypt_and_reply({}, nonce)

    async def get_message(self, msg_json):
        return await parse_wecom_bot_message(msg_json, self.EnCodingAESKey, self.logger)

    async def _handle_message(self, event: wecombotevent.WecomBotEvent):
        """
        处理消息事件。
        """
        try:
            message_id = event.message_id
            if message_id in self.msg_id_map.keys():
                self.msg_id_map[message_id] += 1
                return
            self.msg_id_map[message_id] = 1
            msg_type = event.type
            if msg_type in self._message_handlers:
                for handler in self._message_handlers[msg_type]:
                    await handler(event)
        except Exception:
            print(traceback.format_exc())

    async def push_stream_chunk(self, msg_id: str, content: str, is_final: bool = False) -> bool:
        """将流水线片段推送到 stream 会话。

        Args:
            msg_id: 原始企业微信消息 ID。
            content: 模型产生的片段内容。
            is_final: 是否为最终片段。

        Returns:
            bool: 当成功写入流式队列时返回 True。

        Example:
            在流水线 `reply_message_chunk` 中调用，将增量推送至企业微信。
        """
        # 根据 msg_id 找到对应 stream 会话，如果不存在说明当前消息非流式
        stream_id = self.stream_sessions.get_stream_id_by_msg(msg_id)
        if not stream_id:
            return False

        chunk = StreamChunk(content=content, is_final=is_final)
        await self.stream_sessions.publish(stream_id, chunk)
        if is_final:
            self.stream_sessions.mark_finished(stream_id)
        return True

    async def set_message(self, msg_id: str, content: str):
        """兼容旧逻辑：若无法流式返回则缓存最终结果。

        Args:
            msg_id: 企业微信消息 ID。
            content: 最终回复的文本内容。

        Example:
            在非流式场景下缓存最终结果以备刷新时返回。
        """
        handled = await self.push_stream_chunk(msg_id, content, is_final=True)
        if not handled:
            self.generated_content[msg_id] = content

    def on_message(self, msg_type: str):
        def decorator(func: Callable[[wecombotevent.WecomBotEvent], None]):
            if msg_type not in self._message_handlers:
                self._message_handlers[msg_type] = []
            self._message_handlers[msg_type].append(func)
            return func

        return decorator

    def on_feedback(self):
        def decorator(func: Callable):
            if 'feedback' not in self._message_handlers:
                self._message_handlers['feedback'] = []
            self._message_handlers['feedback'].append(func)
            return func

        return decorator

    async def download_url_to_base64(self, download_url, encoding_aes_key):
        data, _filename = await download_encrypted_file(download_url, encoding_aes_key, self.logger)
        if data:
            return _bytes_to_data_uri(data)
        return None

    async def run_task(self, host: str, port: int, *args, **kwargs):
        """
        启动 Quart 应用。
        """
        await self.app.run_task(host=host, port=port, *args, **kwargs)