geMoldInsight/src/moldinsight/api/batch_router.py

"""
moldinsight/api/batch_router.py — 批量分析端点

- POST /api/batch-upload  批量上传多文件，返回 batch_id + 各 task_id
- GET  /api/batch/{batch_id}  聚合查询批量任务进度
"""
import uuid
from datetime import datetime
from typing import List, Dict, Any, Optional

from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Depends
from sqlalchemy.ext.asyncio import AsyncSession

from shared.database.database import get_db_session
from shared.services.auth_service import get_current_active_user
from shared.models.database import User
from shared.models.schemas import ProcessingStatus, create_task_info
from shared.services.redis_task_manager import redis_task_manager
from shared.utils.file_handler import FileHandler
from shared.utils.logger import get_logger
from moldinsight.services.storage_integration_rustfs import StorageIntegrationService
from moldinsight.services.task_dispatcher import dispatch_processing

logger = get_logger(__name__)

router = APIRouter()

file_handler = FileHandler()

# ─── 批量元数据 Redis key 约定 ──────────────────────────────────────
_BATCH_KEY_PREFIX = "batch:"
_BATCH_TTL = 86400  # 24h

# Redis 不可用时的进程内降级存储（同进程内可查，跨进程/重启不可见）
_batch_meta_memory: Dict[str, dict] = {}


def _batch_redis_key(batch_id: str) -> str:
    return f"{_BATCH_KEY_PREFIX}{batch_id}"


@router.post("/batch-upload")
async def batch_upload(
    files: List[UploadFile] = File(...),
    material: str = Form("ABS"),
    draft_angle: float = Form(2.0),
    shrinkage_rate: float = Form(0.5),
    parting_precision: float = Form(0.1),
    cavity_match: int = Form(95),
    db_session: AsyncSession = Depends(get_db_session),
    current_user: User = Depends(get_current_active_user),
):
    """批量上传多个 STP 文件，每个文件创建独立分析任务，用 batch_id 聚合。"""
    if not files:
        raise HTTPException(400, "请至少上传一个文件")
    if len(files) > 20:
        raise HTTPException(400, "单次批量上传最多 20 个文件")

    process_params = {
        "material": material,
        "draft_angle": float(draft_angle),
        "shrinkage_rate": float(shrinkage_rate),
        "parting_precision": float(parting_precision),
        "cavity_match": int(cavity_match),
    }

    batch_id = str(uuid.uuid4())
    tasks: List[Dict[str, Any]] = []
    storage_service = StorageIntegrationService()

    for file in files:
        # 文件类型检查
        if not file.filename.lower().endswith(('.stp', '.step')):
            tasks.append({
                "filename": file.filename,
                "task_id": None,
                "status": "rejected",
                "error": "不支持的文件类型",
            })
            continue

        task_id = str(uuid.uuid4())
        try:
            file_path, file_size, file_meta = await file_handler.save_uploaded_file(file)

            stp_file = await storage_service.save_stp_file(
                session=db_session,
                file_path=file_path,
                original_filename=file_meta["safe_original_name"],
                user_id=current_user.id,
            )

            await storage_service.create_processing_task(
                db_session, task_id, stp_file.id, parameters=process_params,
            )

            task_info = create_task_info(
                task_id=task_id,
                status=ProcessingStatus.PROCESSING,
                filename=file.filename,
                file_path=str(file_path),
                file_size=file_size,
                upload_time=str(datetime.now()),
            )
            task_info["material"] = material
            task_info["parameters"] = process_params
            task_info["batch_id"] = batch_id
            await redis_task_manager.set_task(task_id, task_info)

            # 调度处理
            dispatch_processing(task_id, str(file_path), stp_file.id, process_params)

            tasks.append({
                "filename": file.filename,
                "task_id": task_id,
                "status": "processing",
                "stp_file_id": stp_file.id,
            })
            logger.info(
                f"[BATCH] batch_id={batch_id} task_id={task_id} "
                f"file={file.filename} user={current_user.username}"
            )
        except Exception as exc:
            logger.warning(f"[BATCH] 文件 {file.filename} 上传失败: {exc}")
            tasks.append({
                "filename": file.filename,
                "task_id": task_id,
                "status": "error",
                "error": str(exc),
            })

    # 将 batch 元数据写入 Redis；Redis 不可用时降级到进程内存储（任务状态本身有内存回退）
    batch_meta = {
        "batch_id": batch_id,
        "user_id": current_user.id,
        "created_at": str(datetime.now()),
        "task_ids": [t["task_id"] for t in tasks if t.get("task_id")],
        "total": len(tasks),
        "params": process_params,
    }
    _save_batch_meta(batch_id, batch_meta)

    return {
        "batch_id": batch_id,
        "total": len(tasks),
        "accepted": sum(1 for t in tasks if t.get("status") != "rejected"),
        "tasks": tasks,
    }


def _save_batch_meta(batch_id: str, batch_meta: dict):
    """批量元数据持久化：优先 Redis（跨进程、带 TTL），降级进程内 dict。"""
    import json as _json

    if redis_task_manager.is_connected:
        try:
            redis_task_manager.redis_client.set(
                _batch_redis_key(batch_id),
                _json.dumps(batch_meta),
                ex=_BATCH_TTL,
            )
            return
        except Exception as exc:
            logger.warning(f"[BATCH] batch 元数据写 Redis 失败，降级内存: {exc}")
    _batch_meta_memory[batch_id] = batch_meta


async def _load_batch_meta(batch_id: str) -> Optional[dict]:
    """读取批量元数据，Redis 优先，内存兜底；不存在返回 None。"""
    import json as _json

    if redis_task_manager.is_connected:
        try:
            raw = await redis_task_manager.redis_client.get(_batch_redis_key(batch_id))
            if raw:
                return _json.loads(raw)
        except Exception as exc:
            logger.warning(f"[BATCH] batch 元数据读 Redis 失败: {exc}")
    return _batch_meta_memory.get(batch_id)


@router.get("/batch/{batch_id}")
async def get_batch_status(
    batch_id: str,
    current_user: User = Depends(get_current_active_user),
):
    """聚合查询批量任务进度"""
    batch_meta = await _load_batch_meta(batch_id)
    if not batch_meta:
        raise HTTPException(404, "批量任务不存在或已过期")

    # 权限检查
    if batch_meta.get("user_id") and batch_meta["user_id"] != current_user.id:
        raise HTTPException(403, "无权访问该批量任务")

    task_ids = batch_meta.get("task_ids", [])
    task_statuses = []
    completed = 0
    failed = 0
    processing = 0

    for tid in task_ids:
        task_data = await redis_task_manager.get_task(tid)
        if not task_data:
            task_statuses.append({"task_id": tid, "status": "unknown"})
            continue
        status = task_data.get("status", "unknown")
        progress = task_data.get("progress", 0)
        filename = task_data.get("filename", "")
        error = task_data.get("error", "")
        html_file = task_data.get("html_file", "")

        if status == ProcessingStatus.COMPLETED:
            completed += 1
        elif status == ProcessingStatus.FAILED:
            failed += 1
        else:
            processing += 1

        task_statuses.append({
            "task_id": tid,
            "status": status,
            "progress": progress,
            "filename": filename,
            "error": error,
            "html_file": html_file,
        })

    total = len(task_ids)
    return {
        "batch_id": batch_id,
        "created_at": batch_meta.get("created_at"),
        "total": total,
        "completed": completed,
        "failed": failed,
        "processing": processing,
        "progress_percent": round((completed + failed) / max(total, 1) * 100, 1),
        "tasks": task_statuses,
    }