chunyu_project/aitool/views/upscale_view.py

"""C-04 图片放大端点：POST /api/aitool/upscale（登录，扣积分，失败回滚）。"""
import logging
import os
import uuid

from adrf.views import APIView
from asgiref.sync import sync_to_async
from django.core.files.storage import default_storage
from django.db import transaction
from django.db.models import F
from django.http import FileResponse, JsonResponse
from rest_framework import status
from rest_framework.parsers import FormParser, MultiPartParser
from rest_framework.permissions import IsAuthenticated

from user.models import FUser, PointTransaction
from utils.async_decorators import async_never_cache_dispatch

from ..models import AIToolRecord
from ..providers import get_provider

logger = logging.getLogger(__name__)

UPSCALE_COST = int(os.environ.get('AITOOL_UPSCALE_COST', '2'))


@async_never_cache_dispatch
class UpscaleView(APIView):
    permission_classes = [IsAuthenticated]
    parser_classes = [MultiPartParser, FormParser]

    async def post(self, request):
        f = request.FILES.get('file') or request.FILES.get('image')
        if not f:
            return JsonResponse({'success': False, 'error': '请上传图片文件'}, status=400)
        try:
            scale = int(request.data.get('scale', 2))
        except (TypeError, ValueError):
            return JsonResponse({'success': False, 'error': 'scale 仅支持 2 或 4'}, status=400)
        if scale not in (2, 4):
            return JsonResponse({'success': False, 'error': 'scale 仅支持 2 或 4'}, status=400)

        try:
            raw = f.read()
        except Exception:
            return JsonResponse({'success': False, 'error': '文件读取失败'}, status=400)
        if not raw:
            return JsonResponse({'success': False, 'error': '空文件'}, status=400)

        # 1) 扣费（事务 + 行锁；余额不足直接拒绝，不建记录）
        def _charge():
            with transaction.atomic():
                u = FUser.objects.select_for_update().get(pk=request.user.pk)
                if u.points < UPSCALE_COST:
                    return None
                u.points = F('points') - UPSCALE_COST
                u.save(update_fields=['points'])
                u.refresh_from_db()
                PointTransaction.objects.create(
                    user=u, transaction_type='spend', currency_type='points',
                    amount=UPSCALE_COST, balance_after=u.points,
                    description=f'AI图片放大 x{scale}',
                )
                return u.points
        try:
            balance = await sync_to_async(_charge)()
        except Exception as e:
            logger.error(f'aitool 扣费失败: {e}')
            return JsonResponse({'success': False, 'error': '扣费失败，请稍后再试'}, status=500)
        if balance is None:
            return JsonResponse(
                {'success': False, 'error': '积分余额不足', 'code': 'INSUFFICIENT_POINTS',
                 'cost': UPSCALE_COST},
                status=402,
            )

        # 2) 处理（失败 → 回滚积分 + 落 failed 记录）
        def _refund(reason: str):
            with transaction.atomic():
                u = FUser.objects.select_for_update().get(pk=request.user.pk)
                u.points = F('points') + UPSCALE_COST
                u.save(update_fields=['points'])
                u.refresh_from_db()
                PointTransaction.objects.create(
                    user=u, transaction_type='earn', currency_type='points',
                    amount=UPSCALE_COST, balance_after=u.points,
                    description=f'AI图片放大失败回滚 x{scale}',
                )
                AIToolRecord.objects.create(
                    user=u, tool='upscale', status='refunded', cost=UPSCALE_COST,
                    scale=scale, input_size=len(raw),
                )
                return u.points
        try:
            out_bytes = await sync_to_async(get_provider().upscale)(raw, scale)
        except ValueError as e:
            await sync_to_async(_refund)(str(e))
            return JsonResponse({'success': False, 'error': str(e)}, status=400)
        except Exception as e:
            logger.error(f'aitool 放大失败: {e}')
            await sync_to_async(_refund)('provider error')
            return JsonResponse({'success': False, 'error': '图片处理失败，已退回积分'}, status=500)

        # 3) 落盘 + 成功记录
        name = f'aitool/upscale/{uuid.uuid4().hex}_{scale}x.png'
        def _save():
            path = default_storage.save(name, __import__('django').core.files.base.ContentFile(out_bytes))
            AIToolRecord.objects.create(
                user_id=request.user.pk, tool='upscale', status='success',
                cost=UPSCALE_COST, scale=scale,
                input_size=len(raw), output_size=len(out_bytes), file_path=path,
            )
            return path
        try:
            saved = await sync_to_async(_save)()
        except Exception as e:
            logger.error(f'aitool 落盘失败: {e}')
            await sync_to_async(_refund)('save error')
            return JsonResponse({'success': False, 'error': '产物保存失败，已退回积分'}, status=500)

        resp = FileResponse(
            __import__('io').BytesIO(out_bytes), content_type='image/png',
            filename=f'upscaled_{scale}x.png', as_attachment=True,
        )
        resp['X-AITool-Cost'] = str(UPSCALE_COST)
        current_balance = await sync_to_async(
            lambda: FUser.objects.get(pk=request.user.pk).points)()
        resp['X-AITool-Balance'] = str(current_balance)
        resp['X-AITool-File'] = f'/media/{saved}'
        return resp


class AIToolRecordsView(APIView):
    """本人调用流水：GET /api/aitool/records（登录）。"""
    permission_classes = [IsAuthenticated]

    async def get(self, request):
        rows = [
            {
                'id': r.id, 'tool': r.tool, 'status': r.status, 'cost': r.cost,
                'scale': r.scale, 'input_size': r.input_size, 'output_size': r.output_size,
                'file_url': f'/media/{r.file_path}' if r.file_path else '',
                'created_at': r.created_at.isoformat() if r.created_at else None,
            }
            async for r in AIToolRecord.objects.filter(user=request.user).order_by('-created_at')[:50]
        ]
        return JsonResponse({'success': True, 'data': rows})