yaoyaoai/app/services/admin_card_key.py

import secrets
import string

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.roles import ROLE_AGENT, ROLE_OEM
from app.models.card_key import CardKey
from app.models.user import User
from app.schemas.admin_card_key import CardKeyCreate, CardKeyUpdate

DEFAULT_PAGE_SIZE = 20
MAX_PAGE_SIZE = 200

_SERIAL_ALPHABET = string.ascii_uppercase + string.digits
_SERIAL_LENGTH = 16


class AdminCardKeyError(Exception):
    def __init__(self, message: str) -> None:
        self.message = message
        super().__init__(message)


def _generate_serial() -> str:
    return "".join(secrets.choice(_SERIAL_ALPHABET) for _ in range(_SERIAL_LENGTH))


async def _unique_serial(db: AsyncSession, preferred: str | None = None) -> str:
    if preferred:
        existing = await db.scalar(
            select(CardKey.id).where(CardKey.serial_number == preferred)
        )
        if existing is not None:
            raise AdminCardKeyError("序列号已存在")
        return preferred

    for _ in range(20):
        candidate = _generate_serial()
        existing = await db.scalar(
            select(CardKey.id).where(CardKey.serial_number == candidate)
        )
        if existing is None:
            return candidate
    raise AdminCardKeyError("生成序列号失败，请重试")


async def _validate_assign_ids(
    db: AsyncSession,
    *,
    oem_id: int | None,
    agent_id: int | None,
) -> None:
    if oem_id is not None and agent_id is not None:
        raise AdminCardKeyError("OEM 与代理只能指定其一")
    if oem_id is not None:
        oem = await db.get(User, oem_id)
        if oem is None or oem.role_id != ROLE_OEM:
            raise AdminCardKeyError("所选 OEM 不存在")
    if agent_id is not None:
        agent = await db.get(User, agent_id)
        if agent is None or agent.role_id != ROLE_AGENT:
            raise AdminCardKeyError("所选代理不存在")


def _list_card_keys_filters(
    *,
    status: str | None = None,
    username: str | None = None,
    oem_id: int | None = None,
    agent_id: int | None = None,
) -> list:
    filters = []
    if status == "unused":
        filters.append(CardKey.activated_at.is_(None))
    elif status == "used":
        filters.append(CardKey.activated_at.is_not(None))
    if username and username.strip():
        filters.append(CardKey.username.ilike(f"%{username.strip()}%"))
    if oem_id is not None:
        filters.append(CardKey.oem_id == oem_id)
    if agent_id is not None:
        filters.append(CardKey.agent_id == agent_id)
    return filters


async def list_card_keys(
    db: AsyncSession,
    *,
    page: int = 1,
    page_size: int = DEFAULT_PAGE_SIZE,
    status: str | None = None,
    username: str | None = None,
    oem_id: int | None = None,
    agent_id: int | None = None,
) -> tuple[list[CardKey], int]:
    page_size = min(max(page_size, 1), MAX_PAGE_SIZE)
    page = max(page, 1)
    offset = (page - 1) * page_size

    filters = _list_card_keys_filters(
        status=status,
        username=username,
        oem_id=oem_id,
        agent_id=agent_id,
    )

    count_stmt = select(func.count()).select_from(CardKey)
    list_stmt = select(CardKey)
    for clause in filters:
        count_stmt = count_stmt.where(clause)
        list_stmt = list_stmt.where(clause)

    total = await db.scalar(count_stmt) or 0
    result = await db.scalars(
        list_stmt.order_by(CardKey.id.desc()).offset(offset).limit(page_size)
    )
    return list(result.all()), total


async def create_card_keys(db: AsyncSession, body: CardKeyCreate) -> list[CardKey]:
    await _validate_assign_ids(db, oem_id=body.oem_id, agent_id=body.agent_id)
    created: list[CardKey] = []

    for index in range(body.count):
        serial = await _unique_serial(
            db,
            body.serial_number if index == 0 else None,
        )
        card = CardKey(
            serial_number=serial,
            duration_days=body.duration_days,
            remark=body.remark,
            oem_id=body.oem_id,
            agent_id=body.agent_id,
        )
        db.add(card)
        created.append(card)

    await db.flush()
    for card in created:
        await db.refresh(card)
    return created


async def update_card_key(
    db: AsyncSession,
    card_id: int,
    body: CardKeyUpdate,
) -> CardKey:
    card = await db.get(CardKey, card_id)
    if card is None:
        raise AdminCardKeyError("卡密不存在")

    if card.activated_at is not None:
        if body.duration_days is not None:
            raise AdminCardKeyError("已激活卡密不能修改时长")
    elif body.duration_days is not None:
        card.duration_days = body.duration_days

    if body.remark is not None:
        card.remark = body.remark

    await db.flush()
    await db.refresh(card)
    return card


async def delete_card_key(db: AsyncSession, card_id: int) -> None:
    card = await db.get(CardKey, card_id)
    if card is None:
        raise AdminCardKeyError("卡密不存在")
    if card.activated_at is not None:
        raise AdminCardKeyError("已激活卡密不能删除")

    await db.delete(card)