geMoldInsight/src/moldinsight/api/export_router.py

# api/export_router.py
"""导出类接口（批次 3 自 advanced_router 拆分，D1）。

导出产物清单（export_artifacts）的合并/校验辅助函数自原文件平移，
行为不变；任务归属校验直接调用 TaskQueryService.ensure_task_access。
"""
from datetime import datetime
from pathlib import Path
from typing import List, Optional
import os
from urllib.parse import quote

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
from sqlalchemy.ext.asyncio import AsyncSession

from shared.services.auth_service import get_current_active_user
from shared.services.redis_task_manager import redis_task_manager
from shared.database.database import get_db_session
from shared.models.identity import User
from moldinsight.services.processing_service import processing_service
from moldinsight.services.task_query_service import TaskQueryService
from moldinsight.services.task_storage_service import TaskStorageService
from moldinsight.core.cad_exporter import CADExporter
from shared.utils.logger import get_logger

logger = get_logger(__name__)

router = APIRouter()
cad_exporter = CADExporter()


class ExportMoldRequest(BaseModel):
    task_id: str
    scheme_id: Optional[str] = None
    formats: List[str] = Field(default_factory=lambda: ["step", "stl"])
    components: List[str] = Field(default_factory=lambda: ["cavity", "core"])


# ---- 导出产物清单辅助（自原 advanced_router 平移） ----

def _get_export_artifacts(task_data: dict) -> dict:
    if not isinstance(task_data, dict):
        return {}
    direct = task_data.get("export_artifacts")
    if isinstance(direct, dict):
        return direct
    parameters = task_data.get("parameters")
    if isinstance(parameters, dict) and isinstance(parameters.get("export_artifacts"), dict):
        return parameters.get("export_artifacts")
    return {}


def _expand_components(components):
    requested = components or ["cavity", "core"]
    if "all" in requested:
        return ["cavity", "core", "parting_surface"]
    return list(dict.fromkeys(requested))


def _augment_export_files(task_id: str, files):
    items = []
    for file in files or []:
        item = dict(file)
        relative_path = item.get("relative_path")
        if not relative_path and item.get("filepath"):
            relative_path = cad_exporter.get_relative_path(item["filepath"])
        if relative_path:
            relative_path = str(relative_path).replace("\\", "/").strip("/")
            item["relative_path"] = relative_path
            item["download_path"] = f"/api/export-download/{quote(relative_path, safe='/')}?task_id={task_id}"
        items.append(item)
    return items


def _merge_export_artifacts(existing: dict, export_result: dict) -> dict:
    merged = dict(existing or {})
    schemes = dict(merged.get("schemes") or {})
    scheme_id = export_result.get("scheme_id") or "default"
    previous = dict(schemes.get(scheme_id) or {})

    file_map = {}
    for file in previous.get("files", []):
        file_map[(file.get("component"), file.get("format"))] = file
    for file in export_result.get("files", []):
        file_map[(file.get("component"), file.get("format"))] = file

    schemes[scheme_id] = {
        "base_filename": export_result.get("base_filename") or previous.get("base_filename"),
        "generated_at": datetime.now().isoformat(),
        "files": sorted(
            file_map.values(),
            key=lambda item: (item.get("component", ""), item.get("format", "")),
        ),
        "errors": export_result.get("errors", []),
        "total_files": len(file_map),
        "total_errors": len(export_result.get("errors", [])),
    }

    merged["version"] = 1
    merged["task_id"] = export_result.get("task_id") or merged.get("task_id")
    merged["generated_at"] = merged.get("generated_at") or datetime.now().isoformat()
    merged["schemes"] = schemes
    return merged


def _select_persisted_files(task_id: str, task_data: dict, scheme_id: str, formats, components):
    artifacts = _get_export_artifacts(task_data)
    scheme_data = (artifacts.get("schemes") or {}).get(scheme_id)
    if not scheme_data:
        return None

    component_list = _expand_components(components)
    format_list = list(dict.fromkeys(formats or ["step", "stl"]))
    expected = {(component, fmt) for component in component_list for fmt in format_list}

    available = []
    available_keys = set()
    for file in scheme_data.get("files", []):
        component = file.get("component")
        fmt = file.get("format")
        if component not in component_list or fmt not in format_list:
            continue
        relative_path = str(file.get("relative_path") or "").replace("\\", "/").strip("/")
        if not relative_path:
            continue
        full_path = os.path.join(cad_exporter.output_dir, relative_path.replace("/", os.sep))
        if not os.path.exists(full_path):
            continue
        available.append(file)
        available_keys.add((component, fmt))

    if expected and not expected.issubset(available_keys):
        return None

    return _augment_export_files(task_id, available)


# ---- 端点 ----

@router.post("/export-mold")
async def export_mold_results(
    body: ExportMoldRequest,
    current_user: User = Depends(get_current_active_user),
    db_session: AsyncSession = Depends(get_db_session),
):
    task_id = body.task_id
    formats = body.formats
    components = body.components

    # 归属校验统一走 TaskQueryService（与 /api/status 共用，含 404/403 语义）
    await TaskQueryService.ensure_task_access(db_session, task_id, current_user.id)
    task_data = await TaskQueryService.get_task_view(db_session, task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")

    resolved_scheme_id = body.scheme_id or task_data.get("best_scheme_id") or "default"
    persisted_files = _select_persisted_files(
        task_id=task_id,
        task_data=task_data,
        scheme_id=resolved_scheme_id,
        formats=formats,
        components=components,
    )
    if persisted_files:
        return {
            "status": "success",
            "data": {
                "base_filename": Path(task_data.get("filename", f"mold_{task_id}")).stem,
                "task_id": task_id,
                "scheme_id": resolved_scheme_id,
                "files": persisted_files,
                "errors": [],
                "total_files": len(persisted_files),
                "total_errors": 0,
                "source": "persisted",
            },
        }

    cavity_shapes = processing_service.get_export_shapes(
        task_id,
        resolved_scheme_id,
    )
    filename = task_data.get("filename", f"mold_{task_id}")

    if not cavity_shapes:
        # 内存 shape 缓存失效（如服务重启）：从持久化的单组件 STEP
        # 现场转换缺失格式，用户无需重新分析
        artifacts = _get_export_artifacts(task_data)
        scheme_data = (artifacts.get("schemes") or {}).get(resolved_scheme_id)
        if scheme_data:
            base_filename = scheme_data.get("base_filename") or Path(filename).stem
            regenerated = await processing_service.regenerate_export_from_persisted(
                task_id=task_id,
                scheme_id=resolved_scheme_id,
                formats=formats,
                components=_expand_components(components),
                base_filename=base_filename,
                scheme_files=scheme_data.get("files", []),
            )
            if regenerated:
                regenerated["files"] = _augment_export_files(
                    task_id, regenerated.get("files", [])
                )
                # 合并进持久化 manifest，后续请求直接命中持久化路径
                merged_artifacts = _merge_export_artifacts(artifacts, regenerated)
                await TaskStorageService().update_task_parameters(
                    db_session,
                    task_id,
                    {"export_artifacts": merged_artifacts},
                )
                # D9：存储方法已不再自行 commit，请求侧显式提交
                await db_session.commit()
                await redis_task_manager.update_task(
                    task_id, {"export_artifacts": merged_artifacts}
                )
                TaskQueryService.invalidate_task_view(task_id)

                return {"status": "success", "data": regenerated}

        raise HTTPException(
            409,
            "导出缓存已失效或任务尚未完成，请重新分析后再导出以保证方案一致性",
        )

    base_filename = Path(filename).stem
    result = cad_exporter.export_mold_results(
        cavity_data=cavity_shapes,
        base_filename=base_filename,
        formats=formats,
        components=components,
        task_id=task_id,
        scheme_id=resolved_scheme_id,
    )
    result["files"] = _augment_export_files(task_id, result.get("files", []))
    result["source"] = "generated"

    merged_artifacts = _merge_export_artifacts(_get_export_artifacts(task_data), result)
    await TaskStorageService().update_task_parameters(
        db_session,
        task_id,
        {"export_artifacts": merged_artifacts},
    )
    # D9：存储方法已不再自行 commit，请求侧显式提交
    await db_session.commit()
    await redis_task_manager.update_task(task_id, {"export_artifacts": merged_artifacts})
    TaskQueryService.invalidate_task_view(task_id)  # parameters 已变更，缓存视图失效

    return {"status": "success", "data": result}


@router.get("/export-download/{filepath:path}")
async def download_export_file(
    filepath: str,
    task_id: str,
    current_user: User = Depends(get_current_active_user),
    db_session: AsyncSession = Depends(get_db_session),
):
    if not task_id:
        raise HTTPException(400, "缺少 task_id")

    # 归属校验统一走 TaskQueryService（与 /api/status 共用，含 404/403 语义）
    await TaskQueryService.ensure_task_access(db_session, task_id, current_user.id)
    task_data = await TaskQueryService.get_task_view(db_session, task_id)
    if not task_data:
        raise HTTPException(404, "任务不存在")

    allowed_paths = set()
    artifacts = _get_export_artifacts(task_data)
    for scheme in (artifacts.get("schemes") or {}).values():
        for file in scheme.get("files", []):
            relative_path = str(file.get("relative_path") or "").replace("\\", "/").strip("/")
            if relative_path:
                allowed_paths.add(relative_path)

    normalized_path = str(filepath or "").replace("\\", "/").strip("/")
    if normalized_path not in allowed_paths:
        raise HTTPException(403, "该文件不在任务允许下载清单中")

    full_path = os.path.join(cad_exporter.output_dir, normalized_path.replace("/", os.sep))
    if not os.path.exists(full_path):
        raise HTTPException(404, "文件不存在")
    if not os.path.abspath(full_path).startswith(os.path.abspath(cad_exporter.output_dir)):
        raise HTTPException(403, "禁止访问")
    media_types = {
        ".step": "application/step", ".stp": "application/step",
        ".iges": "application/iges", ".igs": "application/iges",
        ".stl": "model/stl", ".brep": "application/octet-stream",
    }
    ext = Path(full_path).suffix.lower()
    media_type = media_types.get(ext, "application/octet-stream")
    return FileResponse(full_path, media_type=media_type, filename=os.path.basename(full_path))


@router.get("/export-recommendations")
async def get_export_recommendations(
    target: str = "ug",
    current_user: User = Depends(get_current_active_user),
):
    result = cad_exporter.get_export_recommendations(target)
    return {"status": "success", "data": result}