TelegramBotAdmin/wbb/admin/supervisor.py

from __future__ import annotations

import asyncio
import os
import re
import socket
import sys
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from aiohttp import ClientError, ClientSession, ClientTimeout

from wbb.admin.bot_config import (
    BotConfigError,
    get_bot_profile_secrets,
    telegram_config_status,
)


@dataclass
class BotProcess:
    bot_id: str
    port: int
    process: asyncio.subprocess.Process
    state: str = "starting"
    started_at: str = ""
    error: str = ""
    expected_stop: bool = False
    monitor_task: asyncio.Task[None] | None = None
    readiness_task: asyncio.Task[None] | None = None


class BotSupervisor:
    def __init__(self, config_path: str | Path, *, project_root: str | Path) -> None:
        self.config_path = Path(config_path).resolve()
        self.project_root = Path(project_root).resolve()
        self._processes: dict[str, BotProcess] = {}
        self._lock = asyncio.Lock()

    def runtimes(self) -> dict[str, dict[str, Any]]:
        values: dict[str, dict[str, Any]] = {}
        for bot_id, item in self._processes.items():
            state = item.state
            if item.process.returncode is not None and state not in {"failed", "stopped"}:
                state = "stopped" if item.expected_stop else "failed"
            values[bot_id] = {
                "state": state,
                "pid": item.process.pid,
                "port": item.port,
                "started_at": item.started_at,
                "error": item.error,
            }
        return values

    def endpoint_for(self, bot_id: str) -> str | None:
        item = self._processes.get(str(bot_id))
        if not item or item.state != "running" or item.process.returncode is not None:
            return None
        return f"http://127.0.0.1:{item.port}"

    async def start_enabled(self) -> None:
        status = telegram_config_status(self.config_path)
        for profile in status["bots"]:
            if profile["enabled"] and profile["ready_to_connect"]:
                await self.start(profile["bot_id"])

    async def start(self, bot_id: str) -> dict[str, Any]:
        async with self._lock:
            existing = self._processes.get(str(bot_id))
            if existing and existing.process.returncode is None:
                return self.runtimes()[str(bot_id)]
            profile = get_bot_profile_secrets(self.config_path, bot_id)
            if not profile.get("enabled", True):
                raise BotConfigError("bot_disabled", "该机器人已停用。")
            if not profile.get("api_id") or not profile.get("api_hash"):
                raise BotConfigError(
                    "telegram_api_required",
                    "请先配置 Telegram 应用编号和应用密钥。",
                )
            port = self._available_port()
            environment = self._worker_environment(profile, port)
            process = await asyncio.create_subprocess_exec(
                sys.executable,
                str(self.project_root / "launcher.py"),
                "--worker",
                str(bot_id),
                "--port",
                str(port),
                cwd=self.project_root,
                env=environment,
            )
            item = BotProcess(
                bot_id=str(bot_id),
                port=port,
                process=process,
                started_at=datetime.now(UTC).isoformat().replace("+00:00", "Z"),
            )
            self._processes[str(bot_id)] = item
            item.monitor_task = asyncio.create_task(self._monitor(item))
            item.readiness_task = asyncio.create_task(self._wait_until_ready(item))
            return self.runtimes()[str(bot_id)]

    async def stop(self, bot_id: str) -> dict[str, Any]:
        async with self._lock:
            item = self._processes.get(str(bot_id))
            if not item or item.process.returncode is not None:
                return {"state": "stopped"}
            item.expected_stop = True
            item.state = "stopping"
            item.process.terminate()
        try:
            await asyncio.wait_for(item.process.wait(), timeout=10)
        except TimeoutError:
            item.process.kill()
            await item.process.wait()
        item.state = "stopped"
        return self.runtimes()[str(bot_id)]

    async def restart(self, bot_id: str) -> dict[str, Any]:
        await self.stop(bot_id)
        return await self.start(bot_id)

    async def reconcile(self, bot_id: str) -> dict[str, Any]:
        profile = get_bot_profile_secrets(self.config_path, bot_id)
        if profile.get("enabled", True) and profile.get("api_id") and profile.get("api_hash"):
            return await self.restart(bot_id)
        return await self.stop(bot_id)

    async def close(self) -> None:
        for bot_id in list(self._processes):
            await self.stop(bot_id)

    async def _monitor(self, item: BotProcess) -> None:
        return_code = await item.process.wait()
        if item.readiness_task and not item.readiness_task.done():
            item.readiness_task.cancel()
        if item.expected_stop:
            item.state = "stopped"
            return
        item.state = "failed"
        item.error = f"Bot 工作进程退出，状态码 {return_code}。"

    async def _wait_until_ready(self, item: BotProcess) -> None:
        timeout = ClientTimeout(total=2)
        async with ClientSession(timeout=timeout) as session:
            for _ in range(45):
                if item.process.returncode is not None:
                    return
                try:
                    async with session.get(
                        f"http://127.0.0.1:{item.port}/api/admin/v1/health"
                    ) as response:
                        if response.status == 200:
                            item.state = "running"
                            item.error = ""
                            return
                except (ClientError, TimeoutError):
                    pass
                await asyncio.sleep(1)
        if item.process.returncode is None:
            item.state = "failed"
            item.error = "Bot 已连接，但内部管理接口未在限定时间内就绪。"

    def _worker_environment(self, profile: dict[str, Any], port: int) -> dict[str, str]:
        environment = os.environ.copy()
        environment.update(
            {
                "WBB_ADMIN_BOOTSTRAP": "0",
                "WBB_SUPERVISOR_MODE": "0",
                "WBB_BOT_WORKER": "1",
                "WBB_BOT_PROFILE_ID": str(profile["bot_id"]),
                "WBB_BOT_DATABASE": self._worker_database_name(profile["bot_id"]),
                "BOT_TOKEN": str(profile["bot_token"]),
                "API_ID": str(profile["api_id"]),
                "API_HASH": str(profile["api_hash"]),
                "USERBOT_ENABLED": "0",
                "SESSION_STRING": "",
                "PHONE_NUMBER": "",
                "SUDO_USERS_ID": " ".join(str(item) for item in profile["sudo_users_id"]),
                "LOG_GROUP_ID": str(profile["log_group_id"]),
                "GBAN_LOG_GROUP_ID": str(profile["gban_log_group_id"]),
                "MESSAGE_DUMP_CHAT": str(profile["message_dump_chat"]),
                "ADMIN_WEB_ENABLED": "1",
                "ADMIN_WEB_HOST": "127.0.0.1",
                "ADMIN_WEB_PORT": str(port),
                "BOT_PROFILES_PATH": str(self.config_path),
            }
        )
        return environment

    @staticmethod
    def _worker_database_name(bot_id: str) -> str:
        storage_key = re.sub(r"[^A-Za-z0-9_-]", "_", str(bot_id))[:48]
        if not storage_key:
            raise BotConfigError(
                "invalid_bot_id",
                "机器人编号无法用于创建独立数据空间。",
            )
        return f"wbb_bot_{storage_key}"

    @staticmethod
    def _available_port() -> int:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.bind(("127.0.0.1", 0))
            return int(sock.getsockname()[1])