planManage/app/routers/plans.py

"""Plan + QuotaRule + ModelRoute 管理 API"""

from fastapi import APIRouter, HTTPException

from app import database as db
from app.models import (
    PlanCreate, PlanUpdate, PlanOut, PlanWithRules,
    QuotaRuleCreate, QuotaRuleUpdate, QuotaRuleOut,
    ModelRouteBase, ModelRouteOut,
)

router = APIRouter()


# ── Plan CRUD ─────────────────────────────────────────

@router.get("", response_model=list[PlanWithRules])
async def list_plans(enabled_only: bool = False):
    plans = await db.list_plans(enabled_only=enabled_only)
    result = []
    for p in plans:
        rules = await db.list_quota_rules(p["id"])
        result.append({**p, "quota_rules": rules})
    return result


@router.post("", response_model=dict)
async def create_plan(body: PlanCreate):
    return await db.create_plan(
        name=body.name,
        provider_name=body.provider_name,
        api_key=body.api_key,
        api_base=body.api_base,
        plan_type=body.plan_type,
        supported_models=body.supported_models,
        extra_headers=body.extra_headers,
        extra_config=body.extra_config,
        enabled=body.enabled,
    )


@router.get("/{plan_id}", response_model=PlanWithRules)
async def get_plan(plan_id: str):
    p = await db.get_plan(plan_id)
    if not p:
        raise HTTPException(404, "Plan not found")
    rules = await db.list_quota_rules(plan_id)
    return {**p, "quota_rules": rules}


@router.patch("/{plan_id}")
async def update_plan(plan_id: str, body: PlanUpdate):
    fields = body.model_dump(exclude_unset=True)
    if not fields:
        raise HTTPException(400, "No fields to update")
    ok = await db.update_plan(plan_id, **fields)
    if not ok:
        raise HTTPException(404, "Plan not found or no change")
    return {"ok": True}


@router.delete("/{plan_id}")
async def delete_plan(plan_id: str):
    ok = await db.delete_plan(plan_id)
    if not ok:
        raise HTTPException(404, "Plan not found")
    return {"ok": True}


# ── QuotaRule CRUD ────────────────────────────────────

@router.get("/{plan_id}/rules", response_model=list[QuotaRuleOut])
async def list_rules(plan_id: str):
    return await db.list_quota_rules(plan_id)


@router.post("/{plan_id}/rules", response_model=dict)
async def create_rule(plan_id: str, body: QuotaRuleCreate):
    p = await db.get_plan(plan_id)
    if not p:
        raise HTTPException(404, "Plan not found")
    return await db.create_quota_rule(
        plan_id=plan_id,
        rule_name=body.rule_name,
        quota_total=body.quota_total,
        quota_unit=body.quota_unit,
        refresh_type=body.refresh_type.value,
        interval_hours=body.interval_hours,
        calendar_unit=body.calendar_unit,
        calendar_anchor=body.calendar_anchor,
        enabled=body.enabled,
    )


@router.patch("/rules/{rule_id}")
async def update_rule(rule_id: str, body: QuotaRuleUpdate):
    fields = body.model_dump(exclude_unset=True)
    if "refresh_type" in fields and fields["refresh_type"] is not None:
        fields["refresh_type"] = fields["refresh_type"].value
    ok = await db.update_quota_rule(rule_id, **fields)
    if not ok:
        raise HTTPException(404, "Rule not found or no change")
    return {"ok": True}


@router.delete("/rules/{rule_id}")
async def delete_rule(rule_id: str):
    ok = await db.delete_quota_rule(rule_id)
    if not ok:
        raise HTTPException(404, "Rule not found")
    return {"ok": True}


# ── Model Routes ──────────────────────────────────────

@router.get("/routes/models", response_model=list[ModelRouteOut])
async def get_model_routes():
    return await db.list_model_routes()


@router.post("/routes/models", response_model=dict)
async def set_model_route(body: ModelRouteBase):
    return await db.set_model_route(body.model_name, body.plan_id, body.priority)


@router.delete("/routes/models/{route_id}")
async def delete_model_route(route_id: str):
    ok = await db.delete_model_route(route_id)
    if not ok:
        raise HTTPException(404, "Route not found")
    return {"ok": True}