lark2lark/web/supervisor.py

"""Web 面板后端：FastAPI + Bot 子进程管理 + 日志缓冲 + SSE。

被 main.py 通过 uvicorn 启动。
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import signal
import subprocess
import sys
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import AsyncGenerator, Optional

from fastapi import FastAPI, Request, Response, HTTPException, Depends
from fastapi.responses import (
    HTMLResponse,
    JSONResponse,
    StreamingResponse,
    RedirectResponse,
)
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import yaml

# 确保能 import 项目根的 config
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from config import Config, load_config  # noqa: E402
from web.auth import (  # noqa: E402
    COOKIE_NAME,
    COOKIE_MAX_AGE,
    has_password,
    sign_session,
    verify_password,
    verify_session,
)

logger = logging.getLogger("lark2lark.supervisor")

PROJECT_ROOT = Path(__file__).resolve().parent.parent
BOT_SCRIPT = PROJECT_ROOT / "bot.py"
CONFIG_PATH = PROJECT_ROOT / "config.yaml"
INDEX_HTML = Path(__file__).resolve().parent / "index.html"

LOG_BUFFER_MAX = 500
STATUS_PREFIX = "__STATUS__"
BOT_RESTART_DELAY = 1.0  # 秒，重启间隔


# ---------- 数据模型 ----------

class LoginForm(BaseModel):
    password: str


class ConfigUpdate(BaseModel):
    app_id: Optional[str] = None
    app_secret: Optional[str] = None
    source_chat_ids: Optional[list[str]] = None
    target_chat_ids: Optional[list[str]] = None
    dedup_cache_size: Optional[int] = None
    max_qps: Optional[int] = None
    max_retry: Optional[int] = None
    log_level: Optional[str] = None


@dataclass
class LogEntry:
    seq: int
    ts: float
    level: str
    text: str


class LogBuffer:
    """线程安全环形缓冲 + SSE 订阅广播。"""

    def __init__(self, max_size: int = LOG_BUFFER_MAX) -> None:
        self._lock = threading.Lock()
        self._entries: deque[LogEntry] = deque(maxlen=max_size)
        self._seq = 0
        self._cond = threading.Condition(self._lock)

    def append(self, text: str, level: str = "INFO") -> None:
        with self._cond:
            self._seq += 1
            entry = LogEntry(seq=self._seq, ts=time.time(), level=level, text=text.rstrip("\n"))
            self._entries.append(entry)
            self._cond.notify_all()

    def since(self, seq: int) -> list[LogEntry]:
        with self._lock:
            return [e for e in self._entries if e.seq > seq]

    def all(self) -> list[LogEntry]:
        with self._lock:
            return list(self._entries)

    def wait_for_new(self, last_seq: int, timeout: float = 25.0) -> list[LogEntry]:
        """阻塞等待新日志，超时返回空列表（SSE keepalive）。"""
        with self._cond:
            if not self._cond.wait_for(
                lambda: any(e.seq > last_seq for e in self._entries), timeout=timeout
            ):
                return []
            return [e for e in self._entries if e.seq > last_seq]


# ---------- Bot 子进程管理 ----------

class BotManager:
    """管理 bot.py 子进程：启动、停止、重启、日志采集、状态采集。"""

    def __init__(self, log_buffer: LogBuffer) -> None:
        self._log_buffer = log_buffer
        self._proc: Optional[subprocess.Popen] = None
        self._lock = threading.Lock()
        self._reader_thread: Optional[threading.Thread] = None
        self._watch_thread: Optional[threading.Thread] = None
        self._stop_requested = False
        self._last_status: dict = {}
        self._last_status_ts: float = 0.0
        self._start_time: float = 0.0

    def start(self) -> None:
        with self._lock:
            if self._proc and self._proc.poll() is None:
                return
            self._stop_requested = False
            self._start_time = time.time()
            env = os.environ.copy()
            env["PYTHONUNBUFFERED"] = "1"
            env["PYTHONIOENCODING"] = "utf-8"
            try:
                self._proc = subprocess.Popen(
                    [sys.executable, str(BOT_SCRIPT), "--config", str(CONFIG_PATH)],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=str(PROJECT_ROOT),
                    env=env,
                    encoding="utf-8",
                    errors="replace",
                    bufsize=1,  # 行缓冲
                )
            except Exception as e:
                self._log_buffer.append(f"启动 bot 子进程失败: {e}", "ERROR")
                raise
            self._reader_thread = threading.Thread(
                target=self._read_output, name="bot-reader", daemon=True
            )
            self._reader_thread.start()
            self._watch_thread = threading.Thread(
                target=self._watch, name="bot-watch", daemon=True
            )
            self._watch_thread.start()
            self._log_buffer.append("Bot 子进程已启动", "INFO")

    def stop(self) -> None:
        with self._lock:
            self._stop_requested = True
            proc = self._proc
        if proc and proc.poll() is None:
            self._log_buffer.append("正在停止 Bot 子进程...", "INFO")
            try:
                proc.terminate()
                try:
                    proc.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    proc.kill()
                    proc.wait(timeout=3)
            except Exception as e:
                self._log_buffer.append(f"停止 Bot 异常: {e}", "WARNING")

    def restart(self) -> None:
        self._log_buffer.append("重启 Bot 子进程...", "INFO")
        self.stop()
        time.sleep(BOT_RESTART_DELAY)
        self.start()

    def is_alive(self) -> bool:
        with self._lock:
            return self._proc is not None and self._proc.poll() is None

    def uptime(self) -> int:
        if not self.is_alive():
            return 0
        return int(time.time() - self._start_time)

    @property
    def last_status(self) -> dict:
        return dict(self._last_status) if self._last_status else {}

    def _read_output(self) -> None:
        """读取 bot stdout，识别 __STATUS__ 行，其余写入日志缓冲。"""
        proc = self._proc
        if proc is None or proc.stdout is None:
            return
        try:
            for line in proc.stdout:
                line = line.rstrip("\n")
                if not line:
                    continue
                if line.startswith(STATUS_PREFIX):
                    self._parse_status(line[len(STATUS_PREFIX):])
                    continue
                level = self._detect_level(line)
                self._log_buffer.append(line, level)
        except Exception as e:
            self._log_buffer.append(f"读取 Bot 输出异常: {e}", "ERROR")

    @staticmethod
    def _detect_level(line: str) -> str:
        m = re.search(r"\[(DEBUG|INFO|WARNING|ERROR|CRITICAL)\]", line)
        return m.group(1) if m else "INFO"

    def _parse_status(self, json_str: str) -> None:
        try:
            data = json.loads(json_str)
            self._last_status = data
            self._last_status_ts = time.time()
        except json.JSONDecodeError:
            pass

    def _watch(self) -> None:
        """监控子进程存活，异常退出自动重启。"""
        proc = self._proc
        if proc is None:
            return
        while True:
            rc = proc.wait()
            self._log_buffer.append(f"Bot 子进程退出，返回码={rc}", "WARNING")
            if self._stop_requested:
                break
            self._log_buffer.append(f"{BOT_RESTART_DELAY}s 后自动重启 Bot...", "INFO")
            time.sleep(BOT_RESTART_DELAY)
            try:
                self.start()
            except Exception as e:
                self._log_buffer.append(f"自动重启失败: {e}，10s 后再试", "ERROR")
                time.sleep(10)
                continue
            break  # 新进程已由新 watch 线程接管，本线程退出


# ---------- FastAPI 应用 ----------

app = FastAPI(title="lark2lark 面板", docs_url=None, redoc_url=None)

log_buffer = LogBuffer()
bot_mgr = BotManager(log_buffer)


def get_panel_password() -> str:
    """从当前配置读取面板密码。每次调用都重新加载，支持热更新。"""
    try:
        cfg = load_config()
        return cfg.panel_password
    except Exception:
        return ""


def require_auth(request: Request) -> None:
    """FastAPI 依赖：校验登录 Cookie。"""
    token = request.cookies.get(COOKIE_NAME)
    if not verify_session(token or "", get_panel_password()):
        raise HTTPException(status_code=401, detail="未登录")


@app.on_event("startup")
async def _startup() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    logging.getLogger("lark-oapi").setLevel(logging.WARNING)
    logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
    log_buffer.append("Supervisor 启动", "INFO")
    try:
        bot_mgr.start()
    except Exception as e:
        log_buffer.append(f"启动 Bot 失败: {e}", "ERROR")


@app.on_event("shutdown")
async def _shutdown() -> None:
    log_buffer.append("Supervisor 关闭中", "INFO")
    bot_mgr.stop()


# ---------- 路由 ----------

@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    if not INDEX_HTML.exists():
        return HTMLResponse("<h1>index.html 缺失</h1>", status_code=500)
    return HTMLResponse(INDEX_HTML.read_text(encoding="utf-8"))


@app.post("/api/login")
async def login(form: LoginForm, response: Response):
    pwd = get_panel_password()
    if not has_password(pwd):
        raise HTTPException(status_code=503, detail="面板未设置密码（LARK_PANEL_PASSWORD 环境变量或 config.panel.password）")
    if not verify_password(form.password, pwd):
        raise HTTPException(status_code=401, detail="密码错误")
    token = sign_session(pwd)
    response.set_cookie(
        COOKIE_NAME, token,
        max_age=COOKIE_MAX_AGE,
        httponly=True,
        samesite="lax",
    )
    return {"ok": True}


@app.post("/api/logout")
async def logout(response: Response):
    response.delete_cookie(COOKIE_NAME)
    return {"ok": True}


@app.get("/api/config")
async def get_config(_: None = Depends(require_auth)):
    """返回当前配置，app_secret 脱敏。"""
    if not CONFIG_PATH.exists():
        raise HTTPException(status_code=404, detail="config.yaml 不存在")
    with CONFIG_PATH.open("r", encoding="utf-8") as f:
        data = yaml.safe_load(f) or {}
    # 脱敏
    if data.get("app_secret"):
        data["app_secret"] = "***"
    return data


@app.put("/api/config")
async def update_config(payload: ConfigUpdate, _: None = Depends(require_auth)):
    """更新 config.yaml 并重启 Bot。app_secret 为 '***' 时保留原值。"""
    if not CONFIG_PATH.exists():
        raise HTTPException(status_code=404, detail="config.yaml 不存在")
    with CONFIG_PATH.open("r", encoding="utf-8") as f:
        current = yaml.safe_load(f) or {}

    updates = payload.model_dump(exclude_none=True)
    for key, val in updates.items():
        if key == "app_secret" and val == "***":
            continue  # 保留原值
        current[key] = val

    # 校验：写前快速检查
    try:
        _validate_partial(current)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    with CONFIG_PATH.open("w", encoding="utf-8") as f:
        yaml.safe_dump(current, f, allow_unicode=True, sort_keys=False)

    log_buffer.append("配置已更新，重启 Bot 使其生效", "INFO")
    # 异步重启，避免阻塞 HTTP 响应
    threading.Thread(target=bot_mgr.restart, daemon=True).start()
    return {"ok": True, "note": "配置已保存，Bot 正在重启"}


def _validate_partial(data: dict) -> None:
    """对待写入的配置做基本校验。"""
    if "max_qps" in data:
        qps = data["max_qps"]
        if not isinstance(qps, int) or qps <= 0 or qps > 5:
            raise ValueError(f"max_qps 必须在 (0, 5] 区间，当前 {qps}")
    if "source_chat_ids" in data:
        if not isinstance(data["source_chat_ids"], list) or not data["source_chat_ids"]:
            raise ValueError("source_chat_ids 至少 1 个")
    if "target_chat_ids" in data:
        if not isinstance(data["target_chat_ids"], list) or not data["target_chat_ids"]:
            raise ValueError("target_chat_ids 至少 1 个")


@app.get("/api/status")
async def get_status(_: None = Depends(require_auth)):
    """返回 Bot 运行状态。"""
    status = bot_mgr.last_status
    # 状态超时判定：超过 15s 未更新视为失联
    status_fresh = (time.time() - bot_mgr._last_status_ts) < 15 if bot_mgr._last_status_ts else False
    return {
        "bot_alive": bot_mgr.is_alive(),
        "ws_connected": status.get("ws_connected", False) and status_fresh,
        "uptime": status.get("uptime", 0) if status_fresh else 0,
        "queue_size": status.get("queue_size", 0),
        "dedup_size": status.get("dedup_size", 0),
        "dedup_max": status.get("dedup_max", 0),
        "status_fresh": status_fresh,
        "last_status_ts": bot_mgr._last_status_ts,
    }


@app.get("/api/logs")
async def get_logs(request: Request, since: int = 0, _: None = Depends(require_auth)):
    """SSE 流：实时推送新日志。"""
    async def event_stream() -> AsyncGenerator[bytes, None]:
        last_seq = since
        # 先发送历史
        for entry in log_buffer.since(last_seq):
            last_seq = entry.seq
            yield _format_sse(entry)
        # 再订阅新日志
        while True:
            if await request.is_disconnected():
                break
            new_entries = await asyncio.get_event_loop().run_in_executor(
                None, log_buffer.wait_for_new, last_seq, 25.0
            )
            if not new_entries:
                # keepalive
                yield b": ping\n\n"
                continue
            for entry in new_entries:
                last_seq = entry.seq
                yield _format_sse(entry)

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",  # nginx 不缓冲
        },
    )


def _format_sse(entry: LogEntry) -> bytes:
    data = json.dumps({
        "seq": entry.seq,
        "ts": entry.ts,
        "level": entry.level,
        "text": entry.text,
    }, ensure_ascii=False)
    return f"data: {data}\n\n".encode("utf-8")


@app.get("/api/logs/history")
async def get_logs_history(_: None = Depends(require_auth)):
    """返回全部历史日志（一次性）。"""
    return [
        {"seq": e.seq, "ts": e.ts, "level": e.level, "text": e.text}
        for e in log_buffer.all()
    ]


@app.post("/api/bot/restart")
async def restart_bot(_: None = Depends(require_auth)):
    threading.Thread(target=bot_mgr.restart, daemon=True).start()
    return {"ok": True, "note": "Bot 正在重启"}