ForcePilot/backend/package/yuxi/channel/extensions/telegram/polling.py

from __future__ import annotations

import asyncio
import json
import logging
import os
import time

logger = logging.getLogger(__name__)

def _get_offset_dir() -> str:
    data_dir = os.environ.get(
        "DATA_DIR",
        os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "data"),
    )
    offset_dir = os.path.join(data_dir, "channels", "telegram", "offsets")
    os.makedirs(offset_dir, exist_ok=True)
    return offset_dir


_OFFSET_DIR = _get_offset_dir()

TELEGRAM_API_BASE = "https://api.telegram.org"


async def _answer_callback_quiet(token: str, callback_query_id: str) -> None:
    import httpx

    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(5.0)) as client:
            await client.post(
                f"{TELEGRAM_API_BASE}/bot{token}/answerCallbackQuery",
                json={"callback_query_id": callback_query_id},
            )
    except Exception:
        pass


def _ensure_ptb():
    try:
        import telegram  # noqa: F401
    except ImportError:
        raise ImportError(
            "python-telegram-bot>=21.0 is required for Telegram channel. "
            "Install with: pip install python-telegram-bot[job-queue]"
        )


def _pid_alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False


def _acquire_polling_lease(lease_key: str) -> bool:
    os.makedirs(_OFFSET_DIR, exist_ok=True)
    lease_path = os.path.join(_OFFSET_DIR, f".telegram_polling_lease_{lease_key}")

    if os.path.exists(lease_path):
        try:
            with open(lease_path) as f:
                stale = json.load(f)
            stale_pid = stale.get("pid", 0)
            if stale_pid and _pid_alive(stale_pid):
                logger.warning(
                    "Telegram polling lease held by PID %d for key %s, refusing to start",
                    stale_pid, lease_key,
                )
                return False
            logger.info("Telegram polling lease stale (PID %d dead), re-acquiring", stale_pid)
        except (json.JSONDecodeError, OSError):
            pass

    try:
        with open(lease_path, "w") as f:
            json.dump({"pid": os.getpid(), "key": lease_key, "created_at": int(time.time())}, f)
        logger.info("Telegram polling lease acquired: key=%s pid=%d", lease_key, os.getpid())
        return True
    except OSError:
        logger.error("Failed to write Telegram polling lease file: %s", lease_path)
        return False


def _release_polling_lease(lease_key: str) -> None:
    lease_path = os.path.join(_OFFSET_DIR, f".telegram_polling_lease_{lease_key}")
    try:
        os.remove(lease_path)
        logger.info("Telegram polling lease released: key=%s", lease_key)
    except OSError:
        pass


class TelegramPolling:
    def __init__(self):
        self._running = False
        self._app = None
        self._task: asyncio.Task | None = None
        self._offset: int = 0
        self._last_update_time: float = 0.0
        self._lease_key: str | None = None

    async def start(self, account: dict, queue: asyncio.Queue, abort_event: asyncio.Event) -> None:
        _ensure_ptb()
        from telegram import Update
        from telegram.ext import Application, ContextTypes

        token = account.get("token", "")
        if not token:
            logger.error("Telegram polling: no token configured")
            return

        account_id = account.get("account_id", "default")

        from yuxi.channel.extensions.telegram.config import TelegramConfigAdapter

        lease_key = TelegramConfigAdapter.token_fingerprint(token)
        if not _acquire_polling_lease(lease_key):
            logger.error("Telegram polling lease conflict for account %s, refusing to start", account_id)
            return

        self._lease_key = lease_key
        self._offset = _load_offset(account_id)
        self._last_update_time = time.monotonic()

        self._app = Application.builder().token(token).build()

        async def handle_update(update: Update, context: ContextTypes.DEFAULT_TYPE):
            from yuxi.channel.extensions.telegram.monitor import convert_update_to_unified

            self._last_update_time = time.monotonic()
            update_id = update.update_id

            update_dict = update.to_dict()
            unified = convert_update_to_unified(update_dict, account_id)
            if unified:
                try:
                    queue.put_nowait(unified)
                except asyncio.QueueFull:
                    logger.warning("Telegram message queue full, dropping message")

            if "callback_query" in update_dict:
                cq_id = update_dict["callback_query"].get("id", "")
                if cq_id:
                    try:
                        await _answer_callback_quiet(token, cq_id)
                    except Exception:
                        pass

            _persist_offset(account_id, update_id + 1)

        self._app.add_handler(
            __import__("telegram.ext").ext.MessageHandler(
                __import__("telegram.ext").ext.filters.ALL, handle_update
            )
        )

        self._running = True
        logger.info("Telegram polling started for account %s", account_id)

        try:
            await self._app.initialize()
            await self._app.start()
            await self._app.updater.start_polling(
                allowed_updates=[
                    "message", "edited_message", "callback_query",
                    "my_chat_member", "poll", "poll_answer", "message_reaction",
                ],
            )

            while self._running and not abort_event.is_set():
                if time.monotonic() - self._last_update_time > 120:
                    logger.warning(
                        "Telegram polling watchdog for account %s (no updates in 120s)",
                        account_id,
                    )
                await asyncio.sleep(30)

        except asyncio.CancelledError:
            pass
        except Exception:
            logger.exception("Telegram polling error for account %s", account_id)
        finally:
            await self._stop_app()

    async def stop(self) -> None:
        self._running = False
        await self._stop_app()
        if self._lease_key:
            _release_polling_lease(self._lease_key)
            self._lease_key = None

    async def _stop_app(self) -> None:
        if self._app:
            try:
                if self._app.updater:
                    await self._app.updater.stop()
                await self._app.stop()
                await self._app.shutdown()
            except Exception:
                logger.debug("Error during PTB app shutdown", exc_info=True)
            self._app = None


def _offset_path(account_id: str) -> str:
    os.makedirs(_OFFSET_DIR, exist_ok=True)
    return os.path.join(_OFFSET_DIR, f".telegram_offset_{account_id}")


def _load_offset(account_id: str) -> int:
    try:
        with open(_offset_path(account_id)) as f:
            data = json.load(f)
            return data.get("offset", 0)
    except (FileNotFoundError, json.JSONDecodeError):
        return 0


def _persist_offset(account_id: str, offset: int) -> None:
    try:
        with open(_offset_path(account_id), "w") as f:
            json.dump({"offset": offset, "updated_at": int(time.time())}, f)
    except OSError:
        pass