qihuaiflow/backend/python/app/api/token_usage.py

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.orm import Session
from sqlalchemy import and_, desc, asc, func, extract
from pydantic import BaseModel, Field
from typing import Optional
import uuid
from datetime import datetime, date, timedelta, timezone

from ..database import get_db
from ..models import (
    User, TokenUsageLog, TokenBalance, AIModel,
    gen_uuid
)
from ..auth import get_current_user
from ..middleware.log_service import LogService

router = APIRouter()

# CST timezone (UTC+8)
CST_TZ = timezone(timedelta(hours=8))


class RechargeRequest(BaseModel):
    amount: int = Field(ge=1)
    channel: str = Field(default="simulation")


@router.get("/overview")
def get_overview(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    # 总消耗
    total_used = db.query(func.coalesce(func.sum(TokenUsageLog.total_tokens), 0)).filter(
        TokenUsageLog.user_id == current_user.id
    ).scalar()

    # 今日使用中应用数 - 使用CST时区的"今天"
    now_cst = datetime.now(CST_TZ)
    today_start = now_cst.replace(hour=0, minute=0, second=0, microsecond=0)
    # Convert to UTC for DB comparison since container uses UTC timestamps
    today_start_utc = today_start.astimezone(timezone.utc).replace(tzinfo=None)
    today_apps = db.query(func.count(func.distinct(TokenUsageLog.model_id))).filter(
        and_(
            TokenUsageLog.user_id == current_user.id,
            TokenUsageLog.created_at >= today_start_utc,
        )
    ).scalar() or 0

    # 余额
    balance = db.query(TokenBalance).filter(
        TokenBalance.user_id == current_user.id
    ).first()

    server_time = datetime.now(CST_TZ).isoformat()

    return {
        "code": 0,
        "data": {
            "total_used": total_used or 0,
            "balance": balance.balance if balance else 0,
            "today_apps": today_apps,
            "server_time": server_time,
        },
        "message": "ok",
    }


@router.get("/rankings")
def get_rankings(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    # Token消耗最多的模型
    model_rankings = (
        db.query(
            TokenUsageLog.model_id,
            func.sum(TokenUsageLog.total_tokens).label("total"),
            func.count(TokenUsageLog.id).label("count"),
        )
        .filter(TokenUsageLog.user_id == current_user.id)
        .group_by(TokenUsageLog.model_id)
        .order_by(desc("total"))
        .all()
    )

    models_list = []
    for r in model_rankings:
        model = db.query(AIModel).filter(AIModel.id == r.model_id).first()
        models_list.append({
            "model_id": r.model_id,
            "model_name": model.name if model else "unknown",
            "total_tokens": int(r.total),
            "usage_count": int(r.count),
        })

    # TOP3应用
    app_rankings = (
        db.query(
            TokenUsageLog.conversation_id,
            func.sum(TokenUsageLog.total_tokens).label("total"),
        )
        .filter(TokenUsageLog.user_id == current_user.id)
        .group_by(TokenUsageLog.conversation_id)
        .order_by(desc("total"))
        .limit(3)
        .all()
    )

    apps_list = []
    for r in app_rankings:
        apps_list.append({
            "conversation_id": r.conversation_id,
            "total_tokens": int(r.total),
        })

    return {
        "code": 0,
        "data": {
            "model_rankings": models_list,
            "top_apps": apps_list,
        },
        "message": "ok",
    }


@router.get("/logs")
def get_token_logs(
    start_date: Optional[str] = Query(None, description="开始日期 YYYY-MM-DD"),
    end_date: Optional[str] = Query(None, description="结束日期 YYYY-MM-DD"),
    model_id: Optional[str] = Query(None),
    app_type: Optional[str] = Query(None),
    page: int = Query(default=1, ge=1),
    page_size: int = Query(default=20, ge=1, le=100),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    query = db.query(TokenUsageLog).filter(TokenUsageLog.user_id == current_user.id)

    if start_date:
        try:
            start_dt = datetime.strptime(start_date, "%Y-%m-%d")
            query = query.filter(TokenUsageLog.created_at >= start_dt)
        except ValueError:
            raise HTTPException(status_code=400, detail="开始日期格式错误，请使用YYYY-MM-DD")

    if end_date:
        try:
            end_dt = datetime.strptime(end_date, "%Y-%m-%d") + timedelta(days=1)
            query = query.filter(TokenUsageLog.created_at < end_dt)
        except ValueError:
            raise HTTPException(status_code=400, detail="结束日期格式错误，请使用YYYY-MM-DD")

    if model_id:
        query = query.filter(TokenUsageLog.model_id == model_id)

    total = query.count()
    logs = (
        query.order_by(desc(TokenUsageLog.created_at))
        .offset((page - 1) * page_size)
        .limit(page_size)
        .all()
    )

    result = []
    for log_item in logs:
        model = db.query(AIModel).filter(AIModel.id == log_item.model_id).first()
        result.append({
            "id": log_item.id,
            "model_id": log_item.model_id,
            "model_name": model.name if model else "unknown",
            "prompt_tokens": log_item.prompt_tokens or 0,
            "completion_tokens": log_item.completion_tokens or 0,
            "total_tokens": log_item.total_tokens or 0,
            "conversation_id": log_item.conversation_id,
            "created_at": _format_cst_time(log_item.created_at),
        })

    return {
        "code": 0,
        "data": {
            "total": total,
            "page": page,
            "page_size": page_size,
            "items": result,
        },
        "message": "ok",
    }


@router.get("/balance")
def get_balance(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    balance = db.query(TokenBalance).filter(
        TokenBalance.user_id == current_user.id
    ).first()

    return {
        "code": 0,
        "data": {
            "balance": balance.balance if balance else 0,
            "updated_at": str(balance.updated_at) if balance and balance.updated_at else None,
        },
        "message": "ok",
    }


@router.post("/recharge")
def recharge(
    req: RechargeRequest,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    # 获取或创建余额记录
    balance = db.query(TokenBalance).filter(
        TokenBalance.user_id == current_user.id
    ).first()

    if not balance:
        balance = TokenBalance(
            id=gen_uuid(),
            user_id=current_user.id,
            balance=0,
            updated_at=datetime.now(),
        )
        db.add(balance)
        db.flush()

    old_balance = balance.balance or 0
    balance.balance = old_balance + req.amount
    balance.updated_at = datetime.now()

    # 记录操作日志
    LogService.log_operation(
        user_id=current_user.id, username=current_user.username, action="recharge",
        target="token_balance", target_id=balance.id,
        detail=f"充值 {req.amount} tokens，余额从 {old_balance} 变为 {balance.balance}",
    )

    db.commit()

    return {
        "code": 0,
        "data": {
            "balance": balance.balance,
            "recharged": req.amount,
        },
        "message": "充值成功",
    }


def _format_cst_time(dt_value):
    """Format datetime to CST ISO string.
    DB values are written by cst_now() as CST time (stored naive by MySQL).
    No UTC conversion needed - treat naive values as CST directly.
    """
    if dt_value is None:
        return None
    # DB naive values are already CST (written by cst_now()), just attach tz
    return dt_value.replace(tzinfo=CST_TZ).isoformat()