zTools2/app/controllers/whiteboard_controller.py

"""白板接口：REST（访问/管理）+ WebSocket（实时同步）。

路由：
    GET    /api/wb/{board_id}          公开：访问记事本元数据，不存在则新建（前端 init 用）
    WS     /ws/wb/{board_id}           公开：实时协作 + 心跳
    GET    /api/admin/wb               Basic Auth：管理页列表
    DELETE /api/admin/wb/{board_id}    Basic Auth：删除记事本

HTML 页面 /wb/{id} 与管理页 /wb-admin 由 main.py 直接返回静态文件，
不在此 controller 注册，避免与 REST 同路径冲突。
"""

from __future__ import annotations

import json
import logging
import uuid

from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect
from sqlalchemy.orm import Session

from ..database import get_db
from ..dao.whiteboard_dao import WhiteboardDAO
from ..schemas.whiteboard import WhiteboardListResponse, WhiteboardOut
from ..security import require_docs_auth
from ..services.whiteboard_hub import Connection, get_hub
from ..services.whiteboard_service import WhiteboardService

logger = logging.getLogger("zikai.whiteboard")

router = APIRouter(tags=["whiteboard"])


def _service(db: Session = Depends(get_db)) -> WhiteboardService:
    """REST 路径的 service（纯 DB 操作）。"""
    return WhiteboardService(WhiteboardDAO(db))


# ---------------- 公开 REST ----------------

@router.get(
    "/api/wb/{board_id}",
    response_model=WhiteboardOut,
    summary="访问记事本元数据（不存在则新建）",
    description="前端打开 /wb/{id} 页面后调本接口拉取初始文本；不存在时自动创建空板。",
)
def get_whiteboard(board_id: str, service: WhiteboardService = Depends(_service)) -> WhiteboardOut:
    return service.get_or_create(board_id)


# ---------------- 管理 REST（Basic Auth） ----------------

@router.get(
    "/api/admin/wb",
    response_model=WhiteboardListResponse,
    summary="列出所有记事本（需鉴权）",
    description="供记事本管理页使用：board_id / 创建时间 / 编辑次数 / 上次修改时间。",
)
def list_whiteboards(
    limit: int = Query(100, ge=1, le=500),
    offset: int = Query(0, ge=0),
    service: WhiteboardService = Depends(_service),
    _: str = Depends(require_docs_auth),
) -> WhiteboardListResponse:
    total, items = service.list_all(limit=limit, offset=offset)
    return WhiteboardListResponse(total=total, items=items)


@router.delete(
    "/api/admin/wb/{board_id}",
    summary="删除记事本（需鉴权）",
    description="删 DB 行，并关闭该 board 的所有在线 WebSocket 连接。",
)
async def delete_whiteboard(
    board_id: str,
    service: WhiteboardService = Depends(_service),
    _: str = Depends(require_docs_auth),
) -> dict:
    ok = service.delete(board_id)
    if not ok:
        raise HTTPException(404, "白板不存在")
    # 删除成功后踢出该 board 的所有在线连接（close_board 是 async，须在事件循环中调用）
    await get_hub().close_board(board_id)
    return {"deleted": True}


# ---------------- WebSocket（公开，实时同步 + 心跳） ----------------

@router.websocket("/ws/wb/{board_id}")
async def whiteboard_ws(websocket: WebSocket, board_id: str) -> None:
    """白板实时协作端点（文本记事本）。

    协议（JSON 文本帧）：
        client -> server:
            {"type":"hello","client_id":"..."}   首帧（可选，未带则服务端生成）
            {"type":"ping"}                      心跳，服务端回 pong 并刷新计时
            {"type":"edit","content":"..."}      debounce 后发完整文本，持久化并广播给他人
            {"type":"clear"}                     清空，持久化并广播给所有人
        server -> client:
            {"type":"init","content":"...","version":n,"edit_count":m}
            {"type":"pong"}
            {"type":"update","content":"...","version":n,"client_id":"..."}  文本变更
            {"type":"cleared","client_id":"..."}
            {"type":"error","msg":"..."}
    """
    # 路径层只做最基本校验，详细校验交给 service（service 会查表）
    hub = get_hub()
    # 先 accept，便于对非法 board_id 也回一条 error 再关闭
    await websocket.accept()

    # 读取首帧 hello（或任意帧）拿 client_id
    try:
        first = await websocket.receive_text()
    except WebSocketDisconnect:
        return

    client_id = _extract_client_id(first) or uuid.uuid4().hex[:12]

    # 校验 board_id 并加载白板（不存在则新建）
    try:
        board = _with_db(lambda db: WhiteboardService(WhiteboardDAO(db)).get_or_create(board_id))
    except HTTPException as exc:
        await _safe_send(websocket, {"type": "error", "msg": exc.detail})
        await _safe_close(websocket)
        return

    # 注册连接并下发 init
    conn = Connection(websocket=websocket, board_id=board_id, client_id=client_id)
    ok = await hub.register(conn)
    if not ok:
        await _safe_send(websocket, {"type": "error", "msg": "该记事本在线人数已满"})
        await _safe_close(websocket)
        return
    await _safe_send(websocket, {
        "type": "init",
        "content": board.content,
        "version": board.version,
        "edit_count": board.edit_count,
    })

    # 主循环：收消息 -> 处理 -> 广播
    # 单帧大小上限：与 content 限制对齐（256KB 文本 + JSON 开销，留余量到 512KB）
    MAX_FRAME = 512 * 1024
    try:
        while True:
            raw = await websocket.receive_text()
            if len(raw) > MAX_FRAME:
                await _safe_send(websocket, {"type": "error", "msg": "消息过大"})
                continue
            msg = _parse(raw)
            if msg is None:
                continue
            mtype = msg.get("type")
            if mtype == "ping":
                conn.touch()
                await _safe_send(websocket, {"type": "pong"})
                continue
            # 任何有效业务帧都视为活性证据
            conn.touch()
            if mtype == "edit":
                content = msg.get("content")
                if not isinstance(content, str):
                    await _safe_send(websocket, {"type": "error", "msg": "content 必须是字符串"})
                    continue
                try:
                    out = _with_db(
                        lambda db: WhiteboardService(WhiteboardDAO(db)).update_content(board_id, content)
                    )
                except HTTPException as exc:
                    await _safe_send(websocket, {"type": "error", "msg": exc.detail})
                    continue
                # 广播给他人（发送者本地已更新，不回推）
                await hub.broadcast(
                    board_id,
                    {"type": "update", "content": out.content, "version": out.version, "client_id": client_id},
                    exclude=conn,
                )
            elif mtype == "clear":
                try:
                    _with_db(lambda db: WhiteboardService(WhiteboardDAO(db)).clear(board_id))
                except HTTPException as exc:
                    await _safe_send(websocket, {"type": "error", "msg": exc.detail})
                    continue
                # clear 广播给所有人（含发送者，用于确认）
                await hub.broadcast(
                    board_id, {"type": "cleared", "client_id": client_id}
                )
            else:
                await _safe_send(websocket, {"type": "error", "msg": f"未知消息类型 {mtype}"})
    except WebSocketDisconnect:
        pass
    except Exception as exc:  # pragma: no cover
        logger.warning("白板 WS 异常 board=%s client=%s: %s", board_id, client_id, exc)
    finally:
        await hub.disconnect(conn)


# ---------------- helpers ----------------

def _with_db(fn):
    """在独立 Session 中执行 fn 并返回结果；用完即关。供 WS 路径每帧独立事务使用。"""
    from ..database import get_session_local
    db = get_session_local()()
    try:
        return fn(db)
    finally:
        db.close()

def _extract_client_id(raw: str) -> str | None:
    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        return None
    cid = data.get("client_id") if isinstance(data, dict) else None
    if isinstance(cid, str) and cid:
        return cid
    return None


def _parse(raw: str) -> dict | None:
    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        return None
    return data if isinstance(data, dict) else None


async def _safe_send(ws: WebSocket, msg: dict) -> None:
    try:
        await ws.send_json(msg)
    except Exception:  # pragma: no cover
        pass


async def _safe_close(ws: WebSocket) -> None:
    try:
        await ws.close()
    except Exception:  # pragma: no cover
        pass