server-core/app/im/router.py

# -*- coding: utf-8 -*-
"""IM 对外桥接路由（server-core 作为唯一对外入口）。

- ``/im/ws``：WebSocket 桥接到 im-service（透传 JWT）
- ``/im/internal/task-sync|park-sync``：内部同步（供上层按需触发）
- ``/im/{path:path}``：REST 转发到 im-service ``/api/v1/{path}``（透传 JWT）
"""
from __future__ import annotations

import asyncio
import logging

from fastapi import APIRouter, HTTPException, Request, WebSocket, WebSocketDisconnect

from . import client

log = logging.getLogger("im")

router = APIRouter(prefix="/im", tags=["im"])


def _auth(request: Request) -> str:
    return request.headers.get("Authorization", "")


# ── 内部同步（显式优先于 catch-all） ──────────────────────────

@router.get("/health")
async def im_health():
    try:
        data = await client._request(method="GET", path="/health")
        data["ok"] = True
        return data
    except client.IMError as exc:
        return {"ok": False, "detail": str(exc)}


@router.post("/internal/task-sync")
async def internal_task_sync(request: Request):
    body = await _json_body(request)
    task_id = (body or {}).get("task_id", "")
    if not task_id:
        raise HTTPException(status_code=400, detail="task_id 必填")
    ok = await client.sync_task_group(task_id, (body or {}).get("task_title", ""))
    if not ok:
        raise HTTPException(status_code=502, detail="IM 服务不可用")
    return {"ok": True}


@router.post("/internal/park-sync")
async def internal_park_sync(request: Request):
    body = await _json_body(request)
    park_id = (body or {}).get("park_id", "")
    if not park_id:
        raise HTTPException(status_code=400, detail="park_id 必填")
    ok = await client.sync_park_group(park_id, (body or {}).get("park_name", ""))
    if not ok:
        raise HTTPException(status_code=502, detail="IM 服务不可用")
    return {"ok": True}


@router.post("/internal/enterprise-sync")
async def internal_enterprise_sync(request: Request):
    body = await _json_body(request)
    company_id = (body or {}).get("company_id", "")
    if not company_id:
        raise HTTPException(status_code=400, detail="company_id 必填")
    ok = await client.sync_enterprise_group(company_id, (body or {}).get("company_name", ""))
    if not ok:
        raise HTTPException(status_code=502, detail="IM 服务不可用")
    return {"ok": True}


# ── REST 转发 ────────────────────────────────────────────────

@router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE"])
async def im_forward(path: str, request: Request):
    if path.startswith("internal/"):
        raise HTTPException(status_code=404, detail="unknown internal endpoint")
    body = None
    if request.method in ("POST", "PUT", "PATCH"):
        try:
            body = await request.json()
        except Exception:  # noqa: BLE001
            body = None
    # 透传 query 参数（搜索/分页等）
    params = dict(request.query_params) or None
    try:
        return await client.forward(
            "/" + path,
            auth_header=_auth(request),
            method=request.method,
            body=body if body is not None else None,
            params=params,
        )
    except client.IMError as exc:
        raise HTTPException(status_code=502, detail=str(exc)) from exc


# ── WebSocket 桥接 ───────────────────────────────────────────

@router.websocket("/ws")
async def im_ws_bridge(websocket: WebSocket) -> None:
    import websockets

    token = websocket.query_params.get("token", "")
    await websocket.accept()
    im_ws: websockets.WebSocketClientProtocol | None = None
    try:
        url = f"{client._base()}/ws?token={token}"
        im_ws = await websockets.connect(url, max_size=16 * 1024 * 1024)
    except Exception as exc:  # noqa: BLE001
        log.warning("IM WS 桥接上游连接失败: %s", exc)
        await websocket.close(code=1011, reason="IM service unavailable")
        return

    async def client_to_im() -> None:
        try:
            while True:
                raw = await websocket.receive_text()
                await im_ws.send(raw)
        except (WebSocketDisconnect, Exception):  # noqa: BLE001
            pass

    async def im_to_client() -> None:
        import json as _json
        from .client import resolve_avatar_fields
        try:
            while True:
                raw = await im_ws.recv()
                try:
                    parsed = _json.loads(raw)
                    if isinstance(parsed, (dict, list)):
                        raw = _json.dumps(resolve_avatar_fields(parsed), ensure_ascii=False)
                except (ValueError, TypeError):
                    pass  # 非 JSON 帧（如 ping）原样透传
                await websocket.send_text(raw)
        except (websockets.ConnectionClosed, Exception):  # noqa: BLE001
            pass

    tasks = [
        asyncio.create_task(client_to_im()),
        asyncio.create_task(im_to_client()),
    ]
    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
    for t in pending:
        t.cancel()
    try:
        await websocket.close()
    except Exception:  # noqa: BLE001
        pass
    if im_ws is not None:
        try:
            await im_ws.close()
        except Exception:  # noqa: BLE001
            pass


async def _json_body(request: Request) -> dict | None:
    try:
        return await request.json()
    except Exception:  # noqa: BLE001
        return None