lark2lark/bot.py

"""Bot 子进程入口：启动 WebSocket 长连接 + 限流 Worker。

被 supervisor.py 以子进程方式启动，stdout/stderr 被 supervisor 捕获用于日志和状态采集。
定期向 stdout 输出 __STATUS__{...} JSON 行，supervisor 解析后暴露给面板。

也可独立运行：python bot.py
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import signal
import sys
import threading
import time
from typing import Optional

import lark_oapi as lark

from config import Config, load_config
from dedup import MessageDedup
from forwarder import Forwarder, build_event_handler
from ratelimit import RateLimitedWorker

logger = logging.getLogger("lark2lark")

# 状态行前缀，supervisor 据此识别
STATUS_PREFIX = "__STATUS__"
_STATUS_INTERVAL = 5.0  # 秒

# 全局引用，供状态输出线程读取
_dedup: Optional[MessageDedup] = None
_worker: Optional[RateLimitedWorker] = None
_ws_connected = threading.Event()


def setup_logging(level: str) -> None:
    logging.basicConfig(
        level=getattr(logging, level, logging.INFO),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        stream=sys.stdout,  # 关键：输出到 stdout 供 supervisor 捕获
    )
    logging.getLogger("lark-oapi").setLevel(logging.WARNING)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="飞书群聊自动转发机器人（子进程）")
    p.add_argument("--config", "-c", default=None, help="配置文件路径")
    return p.parse_args()


def start_ws_client(cfg: Config, event_handler, ready: threading.Event) -> threading.Thread:
    """在子线程启动 lark WebSocket 客户端。"""
    def _run() -> None:
        cli = lark.ws.Client(
            cfg.app_id,
            cfg.app_secret,
            event_handler=event_handler,
            log_level=lark.LogLevel.INFO,
        )
        ready.set()
        _ws_connected.set()
        logger.info("WebSocket 客户端启动")
        try:
            cli.start()
        except Exception as e:
            logger.exception("WebSocket 客户端异常退出: %s", e)
        finally:
            _ws_connected.clear()

    t = threading.Thread(target=_run, name="lark-ws", daemon=True)
    t.start()
    return t


def _status_loop(start_time: float, stop_event: threading.Event) -> None:
    """每 5 秒向 stdout 输出 __STATUS__{...} 行，供 supervisor 采集。"""
    while not stop_event.wait(_STATUS_INTERVAL):
        status = {
            "ts": int(time.time()),
            "uptime": int(time.time() - start_time),
            "ws_connected": _ws_connected.is_set(),
            "queue_size": 0,
            "dedup_size": 0,
        }
        if _worker is not None:
            # asyncio.Queue 没有线程安全的 qsize 跨线程读，但底层有 _queue
            try:
                status["queue_size"] = _worker._queue.qsize()  # type: ignore[union-attr]
            except Exception:
                pass
        if _dedup is not None:
            stats = _dedup.stats()
            status["dedup_size"] = stats["size"]
            status["dedup_max"] = stats["max_size"]
        sys.stdout.write(STATUS_PREFIX + json.dumps(status, ensure_ascii=False) + "\n")
        sys.stdout.flush()


async def main_async(cfg: Config) -> None:
    global _dedup, _worker
    loop = asyncio.get_running_loop()

    client = lark.Client.builder().app_id(cfg.app_id).app_secret(cfg.app_secret).build()
    _dedup = MessageDedup(max_size=cfg.dedup_cache_size)
    _worker = RateLimitedWorker(max_qps=cfg.max_qps, max_retry=cfg.max_retry)
    forwarder = Forwarder(cfg, client, _dedup, _worker, loop)
    _worker.set_forward_fn(forwarder.forward)

    stop_event = asyncio.Event()
    worker_task = asyncio.create_task(run_worker(_worker, stop_event))

    event_handler = build_event_handler(forwarder)
    ws_ready = threading.Event()
    ws_thread = start_ws_client(cfg, event_handler, ws_ready)
    ws_ready.wait(timeout=5)

    # 状态输出线程
    start_time = time.time()
    status_stop = threading.Event()
    status_thread = threading.Thread(
        target=_status_loop, args=(start_time, status_stop),
        name="status", daemon=True,
    )
    status_thread.start()

    def _on_signal(*_):
        logger.info("收到退出信号，开始关闭...")
        status_stop.set()
        stop_event.set()

    try:
        signal.signal(signal.SIGINT, _on_signal)
        signal.signal(signal.SIGTERM, _on_signal)
    except (ValueError, AttributeError):
        pass

    await worker_task
    status_stop.set()
    logger.info("Bot 已停止")


async def run_worker(worker: RateLimitedWorker, stop_event: asyncio.Event) -> None:
    task = asyncio.create_task(worker.run())
    await stop_event.wait()
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    setup_logging(cfg.log_level)
    logger.info("Bot 启动：源群 %d 个，目标群 %d 个，QPS=%d",
                len(cfg.source_chat_ids), len(cfg.target_chat_ids), cfg.max_qps)
    try:
        asyncio.run(main_async(cfg))
    except KeyboardInterrupt:
        logger.info("已退出")


if __name__ == "__main__":
    main()