shuidrop_gui/Utils/message_models.py

import json
from dataclasses import dataclass, asdict
from typing import Optional, Dict, Any


@dataclass
class PlatformMessage:
    """规范消息结构体 - 严格按照WebSocket文档v2格式"""
    # 必填字段
    type: str  # 消息类型："message" | "staff_list" | "ping" | "pong" | "transfer" | "connect_success" | "error"

    # 条件必填字段（根据消息类型）
    content: Optional[str] = None  # 消息内容
    msg_type: Optional[str] = None  # 消息实际类型："text" | "image" | "video" | "product_card" | "order_card"
    pin_image: Optional[str] = None  # 用户头像URL
    sender: Optional[Dict] = None  # 发送者信息，必须包含id字段
    store_id: Optional[str] = None  # 店铺ID

    # 可选字段
    receiver: Optional[Dict] = None  # 接收者信息
    data: Optional[Dict] = None  # 扩展数据（如客服列表、验证链接等）
    uuid: Optional[str] = None  # 心跳包UUID
    token: Optional[str] = None  # 认证令牌（心跳包可能需要）

    def __post_init__(self):
        """初始化后处理"""
        # 自动设置msg_type（仅对message类型）
        if self.type == "message" and self.msg_type is None and self.content:
            self.msg_type = self._detect_msg_type()

    def _detect_msg_type(self) -> str:
        """自动检测消息类型"""
        if not self.content:
            return "text"

        content = self.content.lower()

        # 图片检测
        if any(ext in content for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp']):
            return "image"

        # 视频检测
        if any(ext in content for ext in ['.mp4', '.avi', '.mov', '.wmv', '.flv']):
            return "video"

        # 订单卡片检测 - 支持多种格式
        if ("商品id：" in self.content and "订单号：" in self.content) or \
                ("商品ID：" in self.content and "订单号：" in self.content) or \
                ("咨询订单号：" in self.content and "商品ID：" in self.content):
            return "order_card"

        # 商品卡片检测
        if any(keyword in content for keyword in ['goods.html', 'item.html', 'item.jd.com', '商品卡片id：']):
            return "product_card"

        # 默认文本消息
        return "text"

    def to_json(self) -> str:
        """序列化为JSON字符串（用于存储）"""
        return json.dumps(self.to_dict(), ensure_ascii=False)

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典，过滤None值"""
        result = {}
        for key, value in asdict(self).items():
            if value is not None:
                result[key] = value
        return result

    @classmethod
    def from_json(cls, json_str: str):
        """从JSON字符串恢复结构体（用于读取）"""
        data = json.loads(json_str)
        return cls(**data)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]):
        """从字典创建实例"""
        return cls(**data)

    @classmethod
    def create_text_message(cls, content: str, sender_id: str, store_id: str, pin_image: str = None):
        """创建文本消息"""
        return cls(
            type="message",
            content=content,
            msg_type="text",
            pin_image=pin_image,
            sender={"id": sender_id},
            store_id=store_id
        )

    @classmethod
    def create_image_message(cls, image_url: str, sender_id: str, store_id: str, pin_image: str = None):
        """创建图片消息"""
        return cls(
            type="message",
            content=image_url,
            msg_type="image",
            pin_image=pin_image,
            sender={"id": sender_id},
            store_id=store_id
        )

    @classmethod
    def create_video_message(cls, video_url: str, sender_id: str, store_id: str, pin_image: str = None):
        """创建视频消息"""
        return cls(
            type="message",
            content=video_url,
            msg_type="video",
            pin_image=pin_image,
            sender={"id": sender_id},
            store_id=store_id
        )

    @classmethod
    def create_order_card_message(cls, product_id: str, order_number: str, sender_id: str, store_id: str,
                                  pin_image: str = None):
        """创建订单卡片消息"""
        return cls(
            type="message",
            content=f"商品id：{product_id} 订单号：{order_number}",
            msg_type="order_card",
            pin_image=pin_image,
            sender={"id": sender_id},
            store_id=store_id
        )

    @classmethod
    def create_product_card_message(cls, product_url: str, sender_id: str, store_id: str, pin_image: str = None):
        """创建商品卡片消息"""
        return cls(
            type="message",
            content=product_url,
            msg_type="product_card",
            pin_image=pin_image,
            sender={"id": sender_id},
            store_id=store_id
        )

    @classmethod
    def create_staff_list_message(cls, staff_list: list, store_id: str):
        """创建客服列表消息"""
        return cls(
            type="staff_list",
            content="客服列表更新",
            data={"staff_list": staff_list},
            store_id=store_id
        )

    @classmethod
    def create_ping_message(cls, uuid_str: str, token: str = None):
        """创建心跳检测消息"""
        message = cls(
            type="ping",
            uuid=uuid_str
        )
        if token:
            message.token = token
        return message

    @classmethod
    def create_pong_message(cls, uuid_str: str):
        """创建心跳回复消息"""
        return cls(
            type="pong",
            uuid=uuid_str
        )