ForcePilot/backend/package/yuxi/channels/adapters/yuanbao/proto_codec.py

from __future__ import annotations

import struct

from yuxi.utils.logging_config import logger

try:
    from google.protobuf import json_format
    from google.protobuf.message import Message as PbMessage

    HAS_PROTOBUF = True
except ImportError:
    HAS_PROTOBUF = False


class ProtoCodec:
    MAGIC_BYTE = 0xFE
    VERSION = 1

    MSG_TYPE_BIZ = 1
    MSG_TYPE_CONN = 2
    MSG_TYPE_HEARTBEAT = 3
    MSG_TYPE_AUTH = 4

    @staticmethod
    def is_protobuf(data: bytes) -> bool:
        return len(data) > 0 and data[0] == ProtoCodec.MAGIC_BYTE

    @staticmethod
    def encode_biz(payload: dict) -> bytes:
        body = ProtoCodec._encode_json(payload)
        return ProtoCodec._build_frame(ProtoCodec.MSG_TYPE_BIZ, body)

    @staticmethod
    def encode_conn(payload: dict) -> bytes:
        body = ProtoCodec._encode_json(payload)
        return ProtoCodec._build_frame(ProtoCodec.MSG_TYPE_CONN, body)

    @staticmethod
    def encode_heartbeat() -> bytes:
        return ProtoCodec._build_frame(ProtoCodec.MSG_TYPE_HEARTBEAT, b"")

    @staticmethod
    def encode_auth(token: str) -> bytes:
        body = ProtoCodec._encode_json({"token": token})
        return ProtoCodec._build_frame(ProtoCodec.MSG_TYPE_AUTH, body)

    @staticmethod
    def decode(data: bytes) -> dict | None:
        if len(data) < 5:
            return None
        try:
            magic = data[0]
            if magic != ProtoCodec.MAGIC_BYTE:
                return None
            version = data[1]
            msg_type = data[2]
            length = struct.unpack(">H", data[3:5])[0]
            body = data[5 : 5 + length]
            return {
                "version": version,
                "msg_type": ProtoCodec._msg_type_name(msg_type),
                "payload": ProtoCodec._decode_json(body),
            }
        except Exception as e:
            logger.warning(f"[ProtoCodec] Decode error: {e}")
            return None

    @staticmethod
    def encode_pb(message: PbMessage) -> bytes:
        if not HAS_PROTOBUF:
            raise ImportError("google.protobuf is not installed")
        body = message.SerializeToString()
        return ProtoCodec._build_frame(ProtoCodec.MSG_TYPE_BIZ, body)

    @staticmethod
    def decode_pb(data: bytes, message_class: type[PbMessage]) -> PbMessage | None:
        if not HAS_PROTOBUF:
            raise ImportError("google.protobuf is not installed")
        decoded = ProtoCodec.decode(data)
        if decoded is None:
            return None
        payload = decoded.get("payload", {})
        try:
            return json_format.ParseDict(payload, message_class())
        except Exception as e:
            logger.warning(f"[ProtoCodec] PB decode error: {e}")
            return None

    @staticmethod
    def pb_to_frame(message: PbMessage, msg_type: int | None = None) -> bytes:
        if not HAS_PROTOBUF:
            raise ImportError("google.protobuf is not installed")
        msg_type = msg_type or ProtoCodec.MSG_TYPE_BIZ
        body = message.SerializeToString()
        return ProtoCodec._build_frame(msg_type, body)

    @staticmethod
    def frame_to_pb(data: bytes, message_class: type[PbMessage]) -> PbMessage | None:
        if not HAS_PROTOBUF:
            raise ImportError("google.protobuf is not installed")
        if len(data) < 5:
            return None
        try:
            magic = data[0]
            if magic != ProtoCodec.MAGIC_BYTE:
                return None
            length = struct.unpack(">H", data[3:5])[0]
            body = data[5 : 5 + length]
            message = message_class()
            message.ParseFromString(body)
            return message
        except Exception as e:
            logger.warning(f"[ProtoCodec] Frame-to-PB decode error: {e}")
            return None

    @staticmethod
    def _build_frame(msg_type: int, body: bytes) -> bytes:
        if len(body) > 65535:
            raise ValueError(f"Body too large: {len(body)} bytes (max 65535)")
        length = struct.pack(">H", len(body))
        return bytes([ProtoCodec.MAGIC_BYTE, ProtoCodec.VERSION, msg_type]) + length + body

    @staticmethod
    def _encode_json(data: dict) -> bytes:
        import json

        return json.dumps(data, ensure_ascii=False).encode("utf-8")

    @staticmethod
    def _decode_json(data: bytes) -> dict:
        import json

        return json.loads(data.decode("utf-8"))

    @staticmethod
    def _msg_type_name(msg_type: int) -> str:
        names = {
            ProtoCodec.MSG_TYPE_BIZ: "biz",
            ProtoCodec.MSG_TYPE_CONN: "conn",
            ProtoCodec.MSG_TYPE_HEARTBEAT: "heartbeat",
            ProtoCodec.MSG_TYPE_AUTH: "auth",
        }
        return names.get(msg_type, "unknown")


def is_protobuf_message(data: bytes) -> bool:
    return ProtoCodec.is_protobuf(data)