lark2lark/forwarder.py

"""事件处理 + 转发核心：白名单过滤、防死循环、去重、调用 Forward API。

lark.ws.Client 在独立线程中运行事件循环并回调 handler。
本模块的事件回调把任务投递到 asyncio.Queue（跨线程安全），
由 ratelimit.Worker 在主事件循环中匀速消费。
"""
from __future__ import annotations

import asyncio
import logging
import threading
from typing import Optional

import lark_oapi as lark
from lark_oapi.api.im.v1 import (
    ForwardMessageRequest,
    ForwardMessageRequestBody,
)

from config import Config
from dedup import MessageDedup
from ratelimit import ForwardTask, RateLimitedWorker

logger = logging.getLogger(__name__)


class Forwarder:
    """绑定配置、去重器、限流 Worker，并提供事件回调。"""

    def __init__(
        self,
        cfg: Config,
        client: lark.Client,
        dedup: MessageDedup,
        worker: RateLimitedWorker,
        loop: asyncio.AbstractEventLoop,
    ) -> None:
        self._cfg = cfg
        self._client = client
        self._dedup = dedup
        self._worker = worker
        self._loop = loop
        self._source_set = set(cfg.source_chat_ids)

    # ---------- 事件回调（lark SDK 在子线程调用） ----------
    def on_message_receive(self, data: lark.im.v1.P2ImMessageReceiveV1) -> None:
        try:
            self._handle_event(data)
        except Exception as e:
            logger.exception("事件处理异常: %s", e)

    def _handle_event(self, data: lark.im.v1.P2ImMessageReceiveV1) -> None:
        event = data.event
        if event is None:
            return

        message = getattr(event, "message", None)
        sender = getattr(event, "sender", None)
        if message is None:
            return

        message_id: str = message.message_id or ""
        chat_id: str = message.chat_id or ""
        sender_type: str = (sender.sender_type if sender else "") or ""

        # 1) 防死循环：丢弃机器人自己发出的消息
        if sender_type == "bot":
            logger.debug("跳过 bot 消息 message_id=%s", message_id)
            return

        # 2) 白名单：只处理源群
        if chat_id not in self._source_set:
            logger.debug("跳过非白名单群 chat_id=%s", chat_id)
            return

        # 3) 去重
        if not self._dedup.check_and_mark(message_id):
            logger.debug("重复消息已丢弃 message_id=%s", message_id)
            return

        logger.info("接收消息 message_id=%s chat_id=%s type=%s",
                    message_id, chat_id, message.message_type)

        # 4) 投递转发任务（每个目标群一个）
        for target in self._cfg.target_chat_ids:
            if target == chat_id:
                # 源==目标，跳过避免无效转发
                continue
            task = ForwardTask(
                message_id=message_id,
                target_chat_id=target,
                source_chat_id=chat_id,
            )
            # 跨线程安全投递到主事件循环的队列
            asyncio.run_coroutine_threadsafe(
                self._worker.enqueue(task), self._loop
            )

    # ---------- 实际转发（Worker 调用，运行在主事件循环） ----------
    async def forward(self, task: ForwardTask) -> bool:
        """调用飞书 Forward API。返回 True 表示成功。"""
        req = (
            ForwardMessageRequest.builder()
            .message_id(task.message_id)
            .receive_id_type("chat_id")
            .request_body(
                ForwardMessageRequestBody.builder()
                .receive_id(task.target_chat_id)
                .build()
            )
            .build()
        )

        def _call() -> bool:
            try:
                resp = self._client.im.v1.message.forward(req)
            except Exception as e:
                logger.error("Forward API 异常 message_id=%s -> %s: %s",
                             task.message_id, task.target_chat_id, e)
                return False
            if not resp.success():
                logger.error("Forward API 失败 message_id=%s -> %s code=%s msg=%s",
                             task.message_id, task.target_chat_id,
                             resp.code, resp.msg)
                return False
            logger.info("转发成功 message_id=%s -> %s",
                        task.message_id, task.target_chat_id)
            return True

        # lark-oapi 的 Client 是同步阻塞调用，放到默认 executor 执行
        return await asyncio.get_event_loop().run_in_executor(None, _call)


def build_event_handler(forwarder: Forwarder):
    """构造 lark EventDispatcherHandler，注册 im.message.receive_v1。"""
    return (
        lark.EventDispatcherHandler.builder("", "")
        .register_p2_im_message_receive_v1(forwarder.on_message_receive)
        .build()
    )