ems-core/backend/app/api/v1/quota.py

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, and_
from datetime import datetime, timezone
from pydantic import BaseModel
from app.core.database import get_db
from app.core.deps import get_current_user, require_roles
from app.models.quota import EnergyQuota, QuotaUsage
from app.models.user import User

router = APIRouter(prefix="/quota", tags=["配额管理"])


class QuotaCreate(BaseModel):
    name: str
    target_type: str
    target_id: int
    energy_type: str
    period: str
    quota_value: float
    unit: str = "kWh"
    warning_threshold_pct: float = 80
    alert_threshold_pct: float = 95


@router.get("")
async def list_quotas(
    target_type: str | None = None,
    energy_type: str | None = None,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(get_current_user),
):
    """列出所有配额，附带当前使用率"""
    query = select(EnergyQuota).where(EnergyQuota.is_active == True)
    if target_type:
        query = query.where(EnergyQuota.target_type == target_type)
    if energy_type:
        query = query.where(EnergyQuota.energy_type == energy_type)
    query = query.order_by(EnergyQuota.id.desc())
    result = await db.execute(query)
    quotas = result.scalars().all()

    items = []
    for q in quotas:
        # 获取最新使用记录
        usage_result = await db.execute(
            select(QuotaUsage)
            .where(QuotaUsage.quota_id == q.id)
            .order_by(QuotaUsage.calculated_at.desc())
            .limit(1)
        )
        usage = usage_result.scalar_one_or_none()
        items.append({
            **_quota_to_dict(q),
            "current_usage": usage.actual_value if usage else 0,
            "usage_rate_pct": usage.usage_rate_pct if usage else 0,
            "usage_status": usage.status if usage else "normal",
        })
    return items


@router.post("")
async def create_quota(
    data: QuotaCreate,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_roles("admin", "energy_manager")),
):
    quota = EnergyQuota(**data.model_dump(), created_by=user.id)
    db.add(quota)
    await db.flush()
    return _quota_to_dict(quota)


@router.put("/{quota_id}")
async def update_quota(
    quota_id: int,
    data: QuotaCreate,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_roles("admin", "energy_manager")),
):
    result = await db.execute(select(EnergyQuota).where(EnergyQuota.id == quota_id))
    quota = result.scalar_one_or_none()
    if not quota:
        raise HTTPException(status_code=404, detail="配额不存在")
    for k, v in data.model_dump(exclude_unset=True).items():
        setattr(quota, k, v)
    return _quota_to_dict(quota)


@router.delete("/{quota_id}")
async def delete_quota(
    quota_id: int,
    db: AsyncSession = Depends(get_db),
    user: User = Depends(require_roles("admin", "energy_manager")),
):
    result = await db.execute(select(EnergyQuota).where(EnergyQuota.id == quota_id))
    quota = result.scalar_one_or_none()
    if not quota:
        raise HTTPException(status_code=404, detail="配额不存在")
    quota.is_active = False
    return {"message": "已删除"}


@router.get("/usage")
async def quota_usage(
    target_type: str | None = None,
    energy_type: str | None = None,
    period: str | None = None,
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=100),
    db: AsyncSession = Depends(get_db),
    user: User = Depends(get_current_user),
):
    """配额使用统计，支持筛选和分页"""
    query = (
        select(QuotaUsage)
        .join(EnergyQuota, QuotaUsage.quota_id == EnergyQuota.id)
        .where(EnergyQuota.is_active == True)
    )
    if target_type:
        query = query.where(EnergyQuota.target_type == target_type)
    if energy_type:
        query = query.where(EnergyQuota.energy_type == energy_type)
    if period:
        query = query.where(EnergyQuota.period == period)

    count_q = select(func.count()).select_from(query.subquery())
    total = (await db.execute(count_q)).scalar()

    query = query.order_by(QuotaUsage.calculated_at.desc()).offset((page - 1) * page_size).limit(page_size)
    result = await db.execute(query)
    return {
        "total": total,
        "items": [{
            "id": u.id, "quota_id": u.quota_id,
            "period_start": str(u.period_start), "period_end": str(u.period_end),
            "actual_value": u.actual_value, "quota_value": u.quota_value,
            "usage_rate_pct": u.usage_rate_pct, "status": u.status,
            "calculated_at": str(u.calculated_at),
        } for u in result.scalars().all()]
    }


@router.get("/compliance")
async def quota_compliance(
    db: AsyncSession = Depends(get_db),
    user: User = Depends(get_current_user),
):
    """配额合规概览：统计各状态数量"""
    # 每个活跃配额的最新使用记录
    quotas_result = await db.execute(
        select(EnergyQuota).where(EnergyQuota.is_active == True)
    )
    quotas = quotas_result.scalars().all()

    summary = {"total": 0, "normal": 0, "warning": 0, "exceeded": 0}
    details = []

    for q in quotas:
        usage_result = await db.execute(
            select(QuotaUsage)
            .where(QuotaUsage.quota_id == q.id)
            .order_by(QuotaUsage.calculated_at.desc())
            .limit(1)
        )
        usage = usage_result.scalar_one_or_none()
        status = usage.status if usage else "normal"
        summary["total"] += 1
        summary[status] = summary.get(status, 0) + 1
        details.append({
            "quota_id": q.id,
            "name": q.name,
            "target_type": q.target_type,
            "energy_type": q.energy_type,
            "quota_value": q.quota_value,
            "actual_value": usage.actual_value if usage else 0,
            "usage_rate_pct": usage.usage_rate_pct if usage else 0,
            "status": status,
        })

    return {"summary": summary, "details": details}


def _quota_to_dict(q: EnergyQuota) -> dict:
    return {
        "id": q.id, "name": q.name, "target_type": q.target_type,
        "target_id": q.target_id, "energy_type": q.energy_type,
        "period": q.period, "quota_value": q.quota_value, "unit": q.unit,
        "warning_threshold_pct": q.warning_threshold_pct,
        "alert_threshold_pct": q.alert_threshold_pct,
        "is_active": q.is_active,
    }