CardScoreDataServer/app/api/auto_import.py

import os
import json
import asyncio
import aiohttp
from typing import Dict, Any, Tuple, Optional, Union
from fastapi import APIRouter, HTTPException, Request, UploadFile, File, Form

from app.core.config import settings
from app.core.logger import get_logger
from app.utils.scheme import CardType

logger = get_logger(__name__)
router = APIRouter()

INFERENCE_LIMIT = 2
UPLOAD_LIMIT = 4
inference_semaphore = asyncio.Semaphore(INFERENCE_LIMIT)
upload_semaphore = asyncio.Semaphore(UPLOAD_LIMIT)

# 推理调用加 semaphore
async def safe_process_main_image(session, file_bytes, filename, score_type, is_reflect_str):
    async with inference_semaphore:
        return await process_main_image(session, file_bytes, filename, score_type, is_reflect_str)

# 上传调用加 semaphore
async def safe_upload_task(task_func, *args, **kwargs):
    async with upload_semaphore:
        return await task_func(*args, **kwargs)


# --- 内部辅助函数 ---
async def call_api_with_bytes(
        session: aiohttp.ClientSession,
        url: str,
        file_bytes: bytes,
        filename: str,
        params: Dict[str, Any] = None,
        form_fields: Dict[str, Any] = None,
        file_field_name: str = 'file'
) -> Tuple[int, bytes]:
    """通用的文件上传API调用函数 (接收字节数据)"""
    form_data = aiohttp.FormData()

    if form_fields:
        for key, value in form_fields.items():
            form_data.add_field(key, str(value))

    # 直接将内存中的字节数据添加到表单
    form_data.add_field(
        file_field_name,
        file_bytes,
        filename=filename or 'image.jpg',
        content_type='image/jpeg'
    )

    try:
        async with session.post(url, data=form_data, params=params) as response:
            response_content = await response.read()
            if not response.ok:
                logger.error(
                    f"[API Error] {url} -> Status: {response.status}, Msg: {response_content.decode('utf-8', errors='ignore')[:100]}")
            return response.status, response_content
    except Exception as e:
        logger.error(f"[Conn Error] {url} -> {e}")
        raise e


async def process_main_image(
        session: aiohttp.ClientSession,
        file_bytes: bytes,
        filename: str,
        score_type: str,
        is_reflect_card: str
) -> Dict[str, Any]:
    """调用推理服务，处理主图片"""
    logger.info(f"处理主图: {score_type} -> {filename}")
    inference_base_url = settings.SCORE_UPDATE_SERVER_URL

    # 1. 获取转正后的图片
    rectify_url = f"{inference_base_url}/api/card_inference/card_rectify_and_center"
    rectify_status, rectified_image_bytes = await call_api_with_bytes(
        session, url=rectify_url, file_bytes=file_bytes, filename=filename
    )
    if rectify_status >= 300:
        raise HTTPException(status_code=500, detail=f"图片转正失败: {score_type}")

    # 2. 获取分数JSON
    score_url = f"{inference_base_url}/api/card_score/score_inference"
    score_params = {"score_type": score_type, "is_reflect_card": is_reflect_card}

    score_status, score_json_bytes = await call_api_with_bytes(
        session, url=score_url, file_bytes=file_bytes, filename=filename, params=score_params
    )
    if score_status >= 300:
        raise HTTPException(status_code=500, detail=f"推理分数失败: {score_type}")

    return {
        "image_type": score_type,
        "rectified_image": rectified_image_bytes,
        "score_json": json.loads(score_json_bytes)
    }


async def create_card_record(session: aiohttp.ClientSession, base_url: str, card_name: str, cardNo: Optional[str],
                             card_type: CardType) -> int:
    """调用自身服务创建新的卡牌记录"""
    url = f"{base_url}{settings.API_PREFIX}/cards/created"
    params = {'card_name': card_name, 'card_type': card_type.value}
    if cardNo:
        params['cardNo'] = cardNo

    async with session.post(url, params=params) as response:
        if response.status == 201:
            data = await response.json()
            return data.get('id')
        else:
            text = await response.text()
            raise HTTPException(status_code=response.status, detail=f"创建卡牌记录失败: {text}")


async def upload_main_image(session: aiohttp.ClientSession, base_url: str, card_id: int,
                            processed_data: Dict[str, Any]):
    """将处理后的主图和JSON上传到自身服务"""
    image_type = processed_data['image_type']
    url = f"{base_url}{settings.API_PREFIX}/images/insert/{card_id}"

    form_data = aiohttp.FormData()
    form_data.add_field('image_type', image_type)
    form_data.add_field('json_data_str', json.dumps(processed_data['score_json'], ensure_ascii=False))
    form_data.add_field(
        'image',
        processed_data['rectified_image'],
        filename=f'{image_type}_rectified.jpg',
        content_type='image/jpeg'
    )

    async with session.post(url, data=form_data) as response:
        if response.status != 201:
            logger.error(f"[主图上传失败] {image_type} code={response.status}")
            raise HTTPException(status_code=500, detail=f"主图保存失败: {image_type}")


async def upload_gray_image(session: aiohttp.ClientSession, base_url: str, card_id: int, image_type: str,
                            file_bytes: bytes, filename: str):
    """将灰度图源文件上传到自身服务"""
    url = f"{base_url}{settings.API_PREFIX}/images/insert/gray/{card_id}"
    form_fields = {'image_type': image_type}

    status, _ = await call_api_with_bytes(
        session, url=url, file_bytes=file_bytes, filename=filename, form_fields=form_fields, file_field_name='image'
    )
    if status != 201:
        logger.error(f"[灰度图上传失败] {image_type} code={status}")
        raise HTTPException(status_code=500, detail=f"灰度图保存失败: {image_type}")


# --- 暴露的API接口 ---

@router.post("/process_and_import", summary="自动化处理并导入卡牌数据")
async def auto_import_script_api(
        request: Request,
        card_name: str = Form(..., description="卡牌名称"),
        cardNo: Optional[str] = Form(None, description="卡牌编号"),
        card_type: CardType = Form(CardType.pokemon, description="卡牌类型"),
        is_reflect_card: bool = Form(True, description="是否是反光卡"),
        strict_mode: bool = Form(False, description="如果为True，必须提供所有4张主图"),

        front_ring: Union[UploadFile, str, None] = File(None, description="正面环光图文件"),
        front_coaxial: Union[UploadFile, str, None] = File(None, description="正面同轴光图文件"),
        back_ring: Union[UploadFile, str, None] = File(None, description="背面环光图文件"),
        back_coaxial: Union[UploadFile, str, None] = File(None, description="背面同轴光图文件"),
        front_gray: Union[UploadFile, str, None] = File(None, description="正面灰度图文件"),
        back_gray: Union[UploadFile, str, None] = File(None, description="背面灰度图文件")
):
    local_base_url = str(request.base_url).rstrip('/')

    main_inputs = {
        "front_ring": front_ring,
        "front_coaxial": front_coaxial,
        "back_ring": back_ring,
        "back_coaxial": back_coaxial
    }
    gray_inputs = {
        "front_gray": front_gray,
        "back_gray": back_gray
    }

    valid_main_files = {
        k: v for k, v in main_inputs.items()
        if (v is not None) and v.filename
    }
    valid_gray_files = {
        k: v for k, v in gray_inputs.items()
        if (v is not None) and v.filename
    }

    provided_main_count = len(valid_main_files)
    if strict_mode and provided_main_count != 4:
        raise HTTPException(status_code=400, detail=f"严格模式开启，必须提供所有4张主图。")
    if not strict_mode and provided_main_count == 0 and not valid_gray_files:
        raise HTTPException(status_code=400, detail="未提供任何图片文件，无法创建。")

    is_reflect_str = "true" if is_reflect_card else "false"

    async with aiohttp.ClientSession() as session:
        try:
            main_bytes_data = {k: (await v.read(), v.filename) for k, v in valid_main_files.items()}
            gray_bytes_data = {k: (await v.read(), v.filename) for k, v in valid_gray_files.items()}

            logger.info(f"--- 开始自动导入任务: {card_name} ---")
            processed_results = []
            for img_type, (f_bytes, f_name) in main_bytes_data.items():
                if len(f_bytes) == 0:
                    raise HTTPException(status_code=400, detail=f"图片文件 {f_name} 内容为空")

                res = await safe_process_main_image(session, f_bytes, f_name, img_type, is_reflect_str)
                processed_results.append(res)

            card_id = await create_card_record(
                session, local_base_url, card_name, cardNo, card_type
            )
            logger.info(f"卡片记录创建成功，ID: {card_id}")

            # ---------- 修改点：safe_upload_task 调用 + gather ----------
            upload_tasks = []
            for res in processed_results:
                upload_tasks.append(safe_upload_task(upload_main_image, session, local_base_url, card_id, res))

            for img_type, (f_bytes, f_name) in gray_bytes_data.items():
                upload_tasks.append(safe_upload_task(upload_gray_image, session, local_base_url, card_id, img_type, f_bytes, f_name))

            if upload_tasks:
                await asyncio.gather(*upload_tasks)
            # ------------------------------------------------------------

            logger.info(f"--- 自动导入流程结束, Card ID: {card_id} ---")
            return {
                "message": "导入成功",
                "card_id": card_id,
                "card_name": card_name,
                "cardNo": cardNo
            }

        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"[流程终止] 发生异常: {e}")
            raise HTTPException(status_code=500, detail=f"自动化处理异常: {str(e)}")

@router.post("/process_and_import_url", summary="通过URL自动化处理并导入卡牌数据")
async def auto_import_url_script_api(
        request: Request,
        card_name: str = Form(..., description="卡牌名称"),
        cardNo: Optional[str] = Form(None, description="卡牌编号"),
        card_type: CardType = Form(CardType.pokemon, description="卡牌类型"),
        is_reflect_card: bool = Form(True, description="是否是反光卡"),
        strict_mode: bool = Form(False, description="如果为True，必须提供所有4张主图URL"),

        front_ring: Optional[str] = Form(None, description="正面环光图URL"),
        front_coaxial: Optional[str] = Form(None, description="正面同轴光图URL"),
        back_ring: Optional[str] = Form(None, description="背面环光图URL"),
        back_coaxial: Optional[str] = Form(None, description="背面同轴光图URL"),
        front_gray: Optional[str] = Form(None, description="正面灰度图URL"),
        back_gray: Optional[str] = Form(None, description="背面灰度图URL")
):
    local_base_url = str(request.base_url).rstrip('/')

    main_inputs = {
        "front_ring": front_ring, "front_coaxial": front_coaxial,
        "back_ring": back_ring, "back_coaxial": back_coaxial
    }
    gray_inputs = {
        "front_gray": front_gray, "back_gray": back_gray
    }

    valid_main_urls = {k: v for k, v in main_inputs.items() if v and v.strip()}
    valid_gray_urls = {k: v for k, v in gray_inputs.items() if v and v.strip()}

    provided_main_count = len(valid_main_urls)
    if strict_mode and provided_main_count != 4:
        raise HTTPException(status_code=400, detail="严格模式开启，必须提供所有4张主图URL。")
    if not strict_mode and provided_main_count == 0 and not valid_gray_urls:
        raise HTTPException(status_code=400, detail="未提供任何图片URL，无法创建。")

    is_reflect_str = "true" if is_reflect_card else "false"

    async with aiohttp.ClientSession() as session:
        try:
            logger.info(f"--- 开始URL自动导入任务: {card_name} ---")

            async def fetch_image(img_key: str, img_url: str):
                try:
                    async with session.get(img_url) as resp:
                        if resp.status != 200:
                            raise HTTPException(status_code=400, detail=f"下载图片失败: {img_key} -> {resp.status}")
                        file_bytes = await resp.read()
                        filename = img_url.split('/')[-1].split('?')[0]
                        if not filename or '.' not in filename:
                            filename = f"{img_key}.jpg"
                        return img_key, (file_bytes, filename)
                except Exception as e:
                    if isinstance(e, HTTPException): raise e
                    raise HTTPException(status_code=400, detail=f"访问图片URL异常: {img_key} -> {str(e)}")

            fetch_tasks = [fetch_image(k, url) for k, url in valid_main_urls.items()]
            fetch_tasks += [fetch_image(k, url) for k, url in valid_gray_urls.items()]

            downloaded_files = await asyncio.gather(*fetch_tasks)

            main_bytes_data = {}
            gray_bytes_data = {}
            for key, data in downloaded_files:
                if key in valid_main_urls:
                    main_bytes_data[key] = data
                else:
                    gray_bytes_data[key] = data

            processed_results = []
            for img_type, (f_bytes, f_name) in main_bytes_data.items():
                if len(f_bytes) == 0:
                    raise HTTPException(status_code=400, detail=f"图片文件 {f_name} 内容为空")
                res = await safe_process_main_image(session, f_bytes, f_name, img_type, is_reflect_str)
                processed_results.append(res)

            card_id = await create_card_record(session, local_base_url, card_name, cardNo, card_type)
            logger.info(f"URL导入卡片记录创建成功，ID: {card_id}")

            # ---------- 修改点：safe_upload_task 调用 + gather ----------
            upload_tasks = []
            for res in processed_results:
                upload_tasks.append(safe_upload_task(upload_main_image, session, local_base_url, card_id, res))

            for img_type, (f_bytes, f_name) in gray_bytes_data.items():
                upload_tasks.append(safe_upload_task(upload_gray_image, session, local_base_url, card_id, img_type, f_bytes, f_name))

            if upload_tasks:
                await asyncio.gather(*upload_tasks)
            # ------------------------------------------------------------

            logger.info(f"--- URL自动导入流程结束, Card ID: {card_id} ---")
            return {
                "message": "URL导入成功",
                "card_id": card_id,
                "card_name": card_name,
                "cardNo": cardNo
            }

        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"[URL导入流程终止] 发生异常: {e}")
            raise HTTPException(status_code=500, detail=f"自动化处理异常: {str(e)}")