api/shop.py

from __future__ import annotations

from io import BytesIO
import re
import uuid

from fastapi import APIRouter, File, Header, HTTPException, UploadFile
from fastapi.responses import FileResponse
from PIL import Image, ImageOps, UnidentifiedImageError
from pydantic import BaseModel, ConfigDict, Field

from api.support import require_admin, require_identity
from services.config import DATA_DIR, config
from services.auth_service import auth_service
from services.pool_pricing_service import pool_pricing_service
from services.shop_service import shop_service


AVATAR_DIR = DATA_DIR / "avatars"
MAX_AVATAR_BYTES = 5 * 1024 * 1024
ALLOWED_AVATAR_TYPES = {"image/jpeg", "image/png", "image/webp", "image/gif"}
ALLOWED_AVATAR_FORMATS = {"JPEG", "PNG", "WEBP", "GIF"}
AVATAR_FILENAME_PATTERN = re.compile(r"^[a-f0-9]{32}\.(?:webp|gif)$")


class RedeemRequest(BaseModel):
    code: str = ""


class CodeBatchCreateRequest(BaseModel):
    count: int = Field(default=1, ge=1, le=1000)
    credits: int = Field(default=1, ge=1)
    prefix: str = "CS"
    batch: str = ""
    note: str = ""
    validity_days: int | None = Field(default=None, ge=1, le=3650)
    credit_expires_days: int | None = Field(default=None, ge=1, le=3650)
    batch_single_use: bool = False
    code_type: str = "normal"
    daily_gift_date: str = ""
    daily_gift_model_ids: list[str] = Field(default_factory=lambda: ["__account_pool__"])
    package_name: str = ""
    package_duration_days: int = Field(default=1, ge=1, le=3650)
    package_daily_limit: int | None = Field(default=None, ge=0, le=1000000)
    package_weekly_limit: int | None = Field(default=None, ge=0, le=1000000)
    package_monthly_limit: int | None = Field(default=None, ge=0, le=1000000)
    package_total_limit: int | None = Field(default=None, ge=0, le=1000000)
    third_party_image_group_id: str = ""
    third_party_image_model_id: str = ""
    third_party_image_auto_enable: bool = True


class CodeDisableRequest(BaseModel):
    code_id: str = ""


class CodeBatchDeleteRequest(BaseModel):
    code_ids: list[str] = Field(min_length=1, max_length=1000)


class CodeCleanupRequest(BaseModel):
    delete_expired: bool = True
    delete_disabled: bool = False
    archive_redeemed: bool = False
    redeemed_older_than_days: int = Field(default=90, ge=1, le=3650)
    limit: int = Field(default=1000, ge=1, le=20000)
    dry_run: bool = True


class CreditAdjustRequest(BaseModel):
    delta: int
    note: str = ""


class ThirdPartyImageWalletAdjustRequest(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    group_id: str
    model_id: str
    delta: int
    note: str = ""


class ProfileUpdateRequest(BaseModel):
    name: str = ""


def create_router() -> APIRouter:
    router = APIRouter()

    @router.get("/api/me/profile")
    async def get_my_profile(authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        profile = auth_service.get_profile(str(identity.get("id") or ""))
        if profile is None and str(identity.get("role") or "").strip().lower() == "admin":
            profile = {
                "id": str(identity.get("id") or "admin"),
                "name": str(identity.get("name") or "管理员"),
                "role": "admin",
                "created_at": None,
                "account_pool_enabled": True,
                "wechat_openid": "",
                "wechat_bound_at": None,
            }
        if profile is None:
            raise HTTPException(status_code=404, detail={"error": "当前用户不存在，请重新登录"})
        return {"profile": profile}

    @router.post("/api/me/profile")
    async def update_my_profile(body: ProfileUpdateRequest, authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        if str(identity.get("role") or "").strip().lower() != "normal":
            raise HTTPException(status_code=403, detail={"error": "只有普通用户可以设置昵称"})
        try:
            item = auth_service.update_self_name(str(identity.get("id") or ""), body.name)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
        return {"profile": auth_service.get_profile(str(item.get("id") or ""))}

    @router.get("/api/avatars/{filename}")
    async def get_avatar(filename: str):
        if not AVATAR_FILENAME_PATTERN.fullmatch(filename):
            raise HTTPException(status_code=404, detail={"error": "头像不存在"})
        avatar_path = AVATAR_DIR / filename
        if not avatar_path.is_file():
            raise HTTPException(status_code=404, detail={"error": "头像不存在"})
        return FileResponse(
            avatar_path,
            media_type="image/gif" if avatar_path.suffix.lower() == ".gif" else "image/webp",
            headers={"Cache-Control": "public, max-age=31536000, immutable"},
        )

    @router.post("/api/me/avatar")
    async def upload_my_avatar(
        file: UploadFile = File(...),
        authorization: str | None = Header(default=None),
    ):
        identity = require_identity(authorization)
        if str(identity.get("role") or "").strip().lower() != "normal":
            raise HTTPException(status_code=403, detail={"error": "只有普通用户可以修改头像"})
        user_id = str(identity.get("id") or "")
        profile = auth_service.get_profile(user_id)
        if profile is None:
            raise HTTPException(status_code=404, detail={"error": "当前用户不存在"})
        if int(profile.get("avatar_updates_remaining") or 0) <= 0:
            raise HTTPException(status_code=429, detail={"error": "头像每天最多修改 3 次，请明天再试"})
        if str(file.content_type or "").lower() not in ALLOWED_AVATAR_TYPES:
            raise HTTPException(status_code=415, detail={"error": "仅支持 JPG、PNG、WebP 或 GIF 图片"})

        content = await file.read(MAX_AVATAR_BYTES + 1)
        await file.close()
        if not content:
            raise HTTPException(status_code=400, detail={"error": "请选择头像图片"})
        if len(content) > MAX_AVATAR_BYTES:
            raise HTTPException(status_code=413, detail={"error": "头像图片不能超过 5MB"})

        try:
            with Image.open(BytesIO(content)) as source:
                detected_format = str(source.format or "").upper()
                source.verify()
            if detected_format not in ALLOWED_AVATAR_FORMATS:
                raise ValueError("unsupported image format")
            if detected_format == "GIF":
                output_bytes = content
                output_suffix = "gif"
            else:
                with Image.open(BytesIO(content)) as source:
                    avatar = ImageOps.exif_transpose(source)
                    avatar.thumbnail((1024, 1024))
                    if avatar.mode not in {"RGB", "RGBA"}:
                        avatar = avatar.convert("RGBA" if "transparency" in avatar.info else "RGB")
                    output = BytesIO()
                    avatar.save(output, format="WEBP", quality=88, method=6)
                    output_bytes = output.getvalue()
                    output_suffix = "webp"
        except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError) as exc:
            raise HTTPException(status_code=400, detail={"error": "图片文件无效或已损坏"}) from exc

        AVATAR_DIR.mkdir(parents=True, exist_ok=True)
        filename = f"{uuid.uuid4().hex}.{output_suffix}"
        avatar_path = AVATAR_DIR / filename
        avatar_path.write_bytes(output_bytes)
        avatar_url = f"/api/avatars/{filename}"
        try:
            item = auth_service.update_self_avatar(user_id, avatar_url)
        except ValueError as exc:
            avatar_path.unlink(missing_ok=True)
            status_code = 429 if "3 次" in str(exc) else 400
            raise HTTPException(status_code=status_code, detail={"error": str(exc)}) from exc
        return {"profile": auth_service.get_profile(str(item.get("id") or ""))}

    @router.get("/api/me/credits")
    async def get_my_credits(authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        credits = auth_service.credit_summary(identity)
        if str(identity.get("role") or "") == "normal":
            quote = pool_pricing_service.quote(identity, 1)
            credits["pool_unit_price"] = quote["unit_price"]
            credits["pool_pay_percentage"] = quote["pay_percentage"]
            credits["pool_rule_name"] = quote.get("rule_name", "")
        return {"credits": credits}

    @router.get("/api/me/redeem-codes")
    async def get_my_redeem_codes(authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        if str(identity.get("role") or "").strip().lower() != "normal":
            raise HTTPException(status_code=403, detail={"error": "只有普通用户可以查看自己的卡密"})
        return {"items": shop_service.list_user_codes(str(identity.get("id") or ""))}

    @router.get("/api/shop/public")
    async def get_shop_public_config(authorization: str | None = Header(default=None)):
        require_identity(authorization)
        return {"purchase_url": config.shop_purchase_url}

    @router.get("/api/me/api-key")
    async def get_my_api_key(authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        if str(identity.get("role") or "").strip().lower() not in {"normal", "user"}:
            raise HTTPException(status_code=403, detail={"error": "只有微信用户或内部用户可以查看密钥"})
        api_key = auth_service.get_api_key_display(str(identity.get("id") or ""))
        if not api_key:
            raise HTTPException(status_code=404, detail={"error": "当前用户密钥无法回显，请联系管理员重置密钥"})
        return {"api_key": api_key}

    @router.post("/api/redeem")
    async def redeem_code(body: RedeemRequest, authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        try:
            return shop_service.redeem(body.code, identity)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc

    @router.post("/api/redeem/preview")
    async def preview_redeem_code(body: RedeemRequest, authorization: str | None = Header(default=None)):
        identity = require_identity(authorization)
        try:
            return shop_service.preview_code(body.code, identity)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc

    @router.get("/api/shop/codes")
    async def list_codes(authorization: str | None = Header(default=None)):
        require_admin(authorization)
        return {"items": shop_service.list_codes()}

    @router.post("/api/shop/codes")
    async def create_codes(body: CodeBatchCreateRequest, authorization: str | None = Header(default=None)):
        require_admin(authorization)
        try:
            return shop_service.generate_codes(
                count=body.count,
                credits=body.credits,
                prefix=body.prefix,
                batch=body.batch,
                note=body.note,
                validity_days=body.validity_days,
                credit_expires_days=body.credit_expires_days,
                batch_single_use=body.batch_single_use,
                code_type=body.code_type,
                daily_gift_date=body.daily_gift_date,
                daily_gift_model_ids=body.daily_gift_model_ids,
                package_name=body.package_name,
                package_duration_days=body.package_duration_days,
                package_daily_limit=body.package_daily_limit,
                package_weekly_limit=body.package_weekly_limit,
                package_monthly_limit=body.package_monthly_limit,
                package_total_limit=body.package_total_limit,
                third_party_image_group_id=body.third_party_image_group_id,
                third_party_image_model_id=body.third_party_image_model_id,
                third_party_image_auto_enable=body.third_party_image_auto_enable,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc

    @router.post("/api/shop/codes/disable")
    async def disable_code(body: CodeDisableRequest, authorization: str | None = Header(default=None)):
        require_admin(authorization)
        try:
            item = shop_service.disable_code(body.code_id)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
        return {"item": item, "items": shop_service.list_codes()}

    @router.post("/api/shop/codes/enable")
    async def enable_code(body: CodeDisableRequest, authorization: str | None = Header(default=None)):
        require_admin(authorization)
        try:
            item = shop_service.enable_code(body.code_id)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
        return {"item": item, "items": shop_service.list_codes()}

    @router.post("/api/shop/codes/delete")
    async def delete_code(body: CodeDisableRequest, authorization: str | None = Header(default=None)):
        require_admin(authorization)
        try:
            item = shop_service.delete_code(body.code_id)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
        return {"item": item, "items": shop_service.list_codes()}

    @router.post("/api/shop/codes/delete-batch")
    async def delete_codes(body: CodeBatchDeleteRequest, authorization: str | None = Header(default=None)):
        require_admin(authorization)
        try:
            return shop_service.delete_codes(body.code_ids)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc

    @router.post("/api/shop/codes/cleanup")
    async def cleanup_codes(body: CodeCleanupRequest, authorization: str | None = Header(default=None)):
        require_admin(authorization)
        try:
            return shop_service.cleanup_codes(
                delete_expired=body.delete_expired,
                delete_disabled=body.delete_disabled,
                delete_redeemed=body.archive_redeemed,
                redeemed_older_than_days=body.redeemed_older_than_days,
                limit=body.limit,
                dry_run=body.dry_run,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc

    @router.get("/api/shop/ledger")
    async def list_ledger(authorization: str | None = Header(default=None)):
        require_admin(authorization)
        return {"items": shop_service.list_ledger()}

    @router.post("/api/shop/users/{key_id}/credits")
    async def adjust_user_credits(
        key_id: str,
        body: CreditAdjustRequest,
        authorization: str | None = Header(default=None),
    ):
        require_admin(authorization)
        try:
            item = shop_service.adjust_user_credits(key_id, body.delta, note=body.note)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
        return {"item": item, "items": auth_service.list_keys(role="normal")}

    @router.post("/api/shop/users/{key_id}/third-party-image-wallet")
    async def adjust_user_third_party_image_wallet(
        key_id: str,
        body: ThirdPartyImageWalletAdjustRequest,
        authorization: str | None = Header(default=None),
    ):
        require_admin(authorization)
        try:
            item = shop_service.adjust_user_third_party_image_wallet(
                key_id,
                group_id=body.group_id,
                model_id=body.model_id,
                delta=body.delta,
                note=body.note,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
        return {"item": item, "items": auth_service.list_keys(role="normal")}

    return router