taiji-AI-PAD/services/mcp-server/app/routes/heicode_proxy.py

"""
Heicode NewAPI 用户视角元数据透传层（P4）

mcp-server 包装 4 个端点暴露给前端，背后调 Heicode NewAPI（admin token）。
前端只对接 mcp-server 一个域，避免直连 NewAPI 引入的 CORS/审计/限流问题。

依据：
- Docs/Heicode-对接进度与待办.md §2.3.1（Heicode 团队 2026-05-07 决策 ④：透传方案）
- Docs/Heicode-完整调用流程图.md §2.5

端点：
- GET /api/user/heicode/balance      → 当前用户余额 + group + status
- GET /api/user/heicode/models       → 当前用户可用模型清单（按 group 过滤的视图）
- GET /api/user/heicode/usage        → 30 天用量数据
- GET /api/user/heicode/logs         → 最近 50 条调用日志
"""
from __future__ import annotations

import uuid
from typing import Any, Dict, List, Optional

from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from database import get_db
from models import User
from app.auth import require_auth
from app.heicode_client import (
    HeicodeNewAPIClient, HeicodeNewAPIError, get_heicode_client,
)
from app.schemas import SuccessResponse


router = APIRouter(prefix="/api/user/heicode", tags=["Heicode NewAPI 透传"])


# ==================== 共享：取当前 mcp-server 用户的 email ====================

async def _current_user_email(principal: dict, db: AsyncSession) -> str:
    """取当前登录用户的 email（用于解析 heicode 本地 user_id）"""
    user_id = principal.get("user_id") or (principal.get("claims") or {}).get("sub")
    if not user_id:
        raise HTTPException(status_code=401, detail="未登录")

    # 优先 claims.email
    email = (principal.get("claims") or {}).get("email") or principal.get("email")
    if email:
        return str(email).strip().lower()

    # fallback 查 DB
    try:
        uid = uuid.UUID(str(user_id))
    except Exception:
        raise HTTPException(status_code=400, detail="user_id 不是有效 UUID")

    user = await db.get(User, uid)
    if user is None or not user.email:
        raise HTTPException(status_code=404, detail="当前用户无 email，无法关联到 Heicode 账号")
    return user.email.strip().lower()


async def _resolve_heicode_uid(
    client: HeicodeNewAPIClient, email: str,
) -> int:
    """email → heicode user_id（int）；找不到 502"""
    if not client.is_configured():
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail={
                "code": "HEICODE_NEWAPI_NOT_CONFIGURED",
                "message": "Heicode NewAPI 未配置 admin token，请联系运维",
            },
        )
    try:
        uid = await client.resolve_user_id_by_email(email)
    except HeicodeNewAPIError as e:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail={
                "code": "HEICODE_NEWAPI_UPSTREAM_ERROR",
                "message": str(e),
            },
        )
    if uid is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail={
                "code": "HEICODE_USER_NOT_FOUND",
                "message": f"在 Heicode NewAPI 找不到 email='{email}' 的用户。"
                           f"用户需先在 Heicode 完成首次登录建账（from-agnet 流程）。",
            },
        )
    return uid


def _request_id(request: Request) -> str:
    return request.headers.get("X-Request-Id") or str(uuid.uuid4())


# ==================== 4 个透传端点 ====================

@router.get("/balance", response_model=SuccessResponse)
async def get_user_balance(
    request: Request,
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db),
):
    """
    返回当前用户在 Heicode NewAPI 上的：余额、group、status、剩余 quota 等关键字段。
    （NewAPI 的 quota 单位见其文档，通常是积分）
    """
    email = await _current_user_email(principal, db)
    client = get_heicode_client()
    uid = await _resolve_heicode_uid(client, email)
    try:
        info = await client.get_user_info(uid, request_id=_request_id(request))
    except HeicodeNewAPIError as e:
        raise HTTPException(status_code=502, detail={"code": "HEICODE_NEWAPI_UPSTREAM_ERROR", "message": str(e)})

    # 提炼前端常用字段（不要直接 dump 上游全部字段，避免泄露多余信息）
    info = info or {}
    return SuccessResponse(data={
        "heicodeUserId": uid,
        "email": email,
        "username": info.get("username"),
        "displayName": info.get("display_name"),
        "group": info.get("group"),
        "status": info.get("status"),
        "quota": info.get("quota"),
        "usedQuota": info.get("used_quota"),
        "requestCount": info.get("request_count"),
    })


@router.get("/models", response_model=SuccessResponse)
async def get_user_models(
    request: Request,
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db),
):
    """
    返回当前 Heicode NewAPI admin 视野内的全部模型清单。
    前端结合 /balance 的 `group` 字段做用户可见性过滤。
    """
    email = await _current_user_email(principal, db)
    client = get_heicode_client()
    uid = await _resolve_heicode_uid(client, email)
    try:
        # 修订 2026-05-08（§7.11.2）：改调 admin /api/user/{id}/models 而非 /api/models
        models = await client.list_user_models(uid, request_id=_request_id(request))
    except HeicodeNewAPIError as e:
        raise HTTPException(status_code=502, detail={"code": "HEICODE_NEWAPI_UPSTREAM_ERROR", "message": str(e)})

    return SuccessResponse(data={
        "heicodeUserId": uid,
        "email": email,
        "items": models,
        "count": len(models) if isinstance(models, list) else 0,
    })


@router.get("/usage", response_model=SuccessResponse)
async def get_user_usage(
    request: Request,
    days: int = Query(30, ge=1, le=90, description="拉取最近多少天的用量数据"),
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db),
):
    """
    返回当前用户在 Heicode NewAPI 上最近 N 天的用量数据（按日聚合）。
    默认 30 天。
    """
    email = await _current_user_email(principal, db)
    client = get_heicode_client()
    uid = await _resolve_heicode_uid(client, email)
    try:
        items = await client.get_user_quota_dates(
            uid, days=days, request_id=_request_id(request),
        )
    except HeicodeNewAPIError as e:
        raise HTTPException(status_code=502, detail={"code": "HEICODE_NEWAPI_UPSTREAM_ERROR", "message": str(e)})

    return SuccessResponse(data={
        "heicodeUserId": uid,
        "email": email,
        "days": days,
        "items": items,
        "count": len(items) if isinstance(items, list) else 0,
    })


@router.get("/logs", response_model=SuccessResponse)
async def get_user_logs(
    request: Request,
    limit: int = Query(50, ge=1, le=200, description="返回的日志条数"),
    page: int = Query(1, ge=1),
    principal: dict = Depends(require_auth),
    db: AsyncSession = Depends(get_db),
):
    """
    返回当前用户在 Heicode NewAPI 上最近 N 条调用日志。
    默认最近 50 条。
    """
    email = await _current_user_email(principal, db)
    client = get_heicode_client()
    uid = await _resolve_heicode_uid(client, email)
    try:
        items = await client.get_user_logs(
            uid, page=page, page_size=limit, request_id=_request_id(request),
        )
    except HeicodeNewAPIError as e:
        raise HTTPException(status_code=502, detail={"code": "HEICODE_NEWAPI_UPSTREAM_ERROR", "message": str(e)})

    return SuccessResponse(data={
        "heicodeUserId": uid,
        "email": email,
        "page": page,
        "limit": limit,
        "items": items,
        "count": len(items) if isinstance(items, list) else 0,
    })