geMoldInsight/src/services/processing_service.py

# services/processing_service.py
"""STP 文件处理流程编排器 — 协调解析、网格生成、型腔生成、保存、验证"""

import asyncio
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any

from sqlalchemy.ext.asyncio import AsyncSession

from core.stp_parser import STPParser
from core.geometry_analyzer import GeometryAnalyzer
from core.mold_generator import MoldCavityGenerator
from core.aluminum_foam_mold import AluminumFoamMoldGenerator
from core.mold_quality_inspector import AluminumFoamMoldQualityInspector
from core.mesh_generator import MeshGenerator
from core.multi_scheme_planner import MultiSchemeMoldPlanner
from core.cad_exporter import CADExporter
from services.storage_integration_rustfs import StorageIntegrationService
from services.redis_task_manager import redis_task_manager
from services.material_service import MaterialService
from services.calculation_service import CalculationService
from services.llm_service import llm_service
from models.schemas import ProcessingStatus
from database.database import db_manager
from utils.html_generator import HTMLGenerator
from utils.logger import get_logger

logger = get_logger(__name__)


class ProcessingService:
    """核心处理流程编排 — 协调 STP 解析、网格、型腔、计算、保存、验证"""

    def __init__(self):
        self.stp_parser = STPParser()
        self.geometry_analyzer = GeometryAnalyzer()
        self.mold_generator = MoldCavityGenerator(shrinkage_rate=0.005)
        self.aluminum_foam_generator = AluminumFoamMoldGenerator(shrinkage_rate=0.015, draft_angle=3.0)
        self.mold_quality_inspector = AluminumFoamMoldQualityInspector()
        self.mesh_generator = MeshGenerator(quality="medium")
        self.html_generator = HTMLGenerator()
        self.storage_service = StorageIntegrationService()
        self.multi_scheme_planner = MultiSchemeMoldPlanner()
        self.cad_exporter = CADExporter()
        self._export_shapes_cache: Dict[str, Dict[str, Dict[str, Any]]] = {}

    # ─── 对外入口 ───

    async def process_file_with_storage(
        self,
        task_id: str,
        file_path: str,
        stp_file_id: int,
        process_params: Optional[Dict[str, Any]] = None,
    ):
        """处理文件的后台任务 — 使用独立数据库会话"""

        # 创建独立的数据库会话，避免请求范围会话关闭
        async with db_manager.session() as db_session:
            try:
                logger.info(f"开始处理文件并生成模具型腔: {file_path}")

                # 设置处理超时（5分钟）
                timeout_seconds = 300

                try:
                    await asyncio.wait_for(
                        self.process_file_core(
                            task_id, file_path, stp_file_id, db_session, process_params
                        ),
                        timeout_seconds,
                    )
                except asyncio.TimeoutError:
                    logger.error(f"处理超时: {task_id}")
                    raise Exception(f"处理超时，超过{timeout_seconds}秒未完成")

            except Exception as e:
                logger.error(f"模具型腔生成失败: {e}")

                await self.storage_service.update_stp_file_status(db_session, stp_file_id, "failed")
                await self.storage_service.update_task_status(
                    db_session, task_id, "failed", error_message=str(e)
                )

                # 安全更新 Redis 任务状态
                task = await redis_task_manager.get_task(task_id)
                if task:
                    await redis_task_manager.update_task(task_id, {
                        "status": ProcessingStatus.FAILED,
                        "error": str(e),
                        "completed_at": str(datetime.now()),
                    })

    async def process_file_core(
        self,
        task_id: str,
        file_path: str,
        stp_file_id: int,
        db_session: AsyncSession,
        process_params: Optional[Dict[str, Any]] = None,
    ):
        """核心处理逻辑"""

        try:
            logger.info(f"开始处理文件并生成模具型腔: {file_path}")
            process_params = self._normalize_process_params(process_params)
            stage_timings: Dict[str, float] = {}

            # 1. 解析STP文件
            await self.storage_service.update_task_status(
                db_session, task_id, "processing", 20, "解析STP文件"
            )

            stage_started = time.perf_counter()
            shape = self.stp_parser.load_step_file(Path(file_path))
            geometry_data = self.stp_parser.analyze_geometry(shape)
            stage_timings["parse_stp"] = round(time.perf_counter() - stage_started, 3)

            # 2. 生成网格数据并持久化
            await self.storage_service.update_task_status(
                db_session, task_id, "processing", 30, "生成网格数据"
            )

            stage_started = time.perf_counter()
            mesh_result = await self._step_generate_mesh(
                shape, geometry_data, file_path, db_session, stp_file_id, task_id
            )
            stage_timings["generate_mesh"] = round(time.perf_counter() - stage_started, 3)

            # 3. 生成模具型腔
            await self.storage_service.update_task_status(
                db_session, task_id, "processing", 40, "生成模具型腔"
            )

            # 材料属性 — 通过 MaterialService 集中管理
            requested_material = MaterialService.resolve_material(process_params["material"])
            selected_material = dict(MaterialService.get_material(requested_material))
            selected_material["shrinkage"] = process_params["shrinkage_rate"] / 100.0
            is_foam_material = MaterialService.is_foam_material(requested_material)

            stage_started = time.perf_counter()
            plan_result = await self._step_generate_cavity(
                shape, selected_material, is_foam_material, process_params
            )
            stage_timings["generate_cavity"] = round(time.perf_counter() - stage_started, 3)
            export_shapes = {}
            export_artifacts = None
            if plan_result:
                export_shapes = plan_result.pop("_export_shapes", {}) or {}
                if export_shapes:
                    self._cache_export_shapes(task_id, export_shapes)
                    export_artifacts = self._persist_step_exports(
                        task_id=task_id,
                        original_filename=Path(file_path).name,
                        export_shapes=export_shapes,
                    )

            # 4. 生成详细JSON数据 — 委托 CalculationService
            await self.storage_service.update_task_status(
                db_session, task_id, "processing", 60, "生成型腔详细数据"
            )

            stage_started = time.perf_counter()
            detailed_cavity_json = CalculationService.build_plan_result(
                geometry_data=geometry_data,
                material=selected_material,
                file_path=str(file_path),
                plan_result=plan_result,
            )
            stage_timings["build_plan_result"] = round(time.perf_counter() - stage_started, 3)

            best_scheme = CalculationService.get_best_scheme(detailed_cavity_json)
            best_cavity_data = best_scheme.get("cavity_data", {}) if best_scheme else {}
            best_key_info = best_scheme.get("key_info", {}) if best_scheme else {}

            if best_cavity_data.get("mold_cavities"):
                cavity_geometry = best_cavity_data["mold_cavities"].get("cavity", {})
                logger.info(
                    f"推荐方案型腔数据已合并: cavity {cavity_geometry.get('vertex_count', 0)} 顶点"
                )

            # 5. 生成关键信息
            cavity_key_info = best_key_info

            # 6. 保存几何数据到数据库
            await self.storage_service.update_task_status(
                db_session, task_id, "processing", 70, "保存几何数据"
            )

            stage_started = time.perf_counter()
            await self.storage_service.save_geometry_data(
                db_session,
                stp_file_id,
                geometry_data,
                geometry_data.get("analysis_method", "mold_cavity"),
            )

            # 7. 生成HTML可视化
            await self.storage_service.update_task_status(
                db_session, task_id, "processing", 85, "生成可视化报告"
            )

            pointcloud_data = None
            lod_data = None
            if mesh_result:
                pointcloud_data = {
                    "points": mesh_result.get("points", []),
                    "normals": mesh_result.get("normals", []),
                    "vertices": mesh_result.get("vertices", []),
                    "faces": mesh_result.get("faces", []),
                    "point_count": mesh_result.get("point_count", 0),
                    "vertex_count": mesh_result.get("vertex_count", 0),
                    "face_count": mesh_result.get("face_count", 0),
                }

            # 生成多级LOD数据（用于前端按距离切换精度）
            try:
                lod_result = self.mesh_generator.generate_multi_lod_mesh(shape)
                if lod_result and lod_result.get("lods"):
                    lod_data = lod_result
                    logger.info(f"LOD数据生成成功: {len(lod_result['lods'])} 级 (面数: {[lod_result['lods'][k]['face_count'] for k in sorted(lod_result['lods'].keys())]})")
            except Exception as lod_err:
                logger.warning(f"LOD数据生成失败，使用单级精度: {lod_err}")

            detailed_cavity_json = await self._attach_scheme_previews(
                detailed_cavity_json=detailed_cavity_json,
                geometry_data=geometry_data,
                stp_filename=Path(file_path).name,
                pointcloud_data=pointcloud_data,
                lod_data=lod_data,
            )

            best_scheme = CalculationService.get_best_scheme(detailed_cavity_json)
            best_cavity_data = best_scheme.get("cavity_data", {}) if best_scheme else best_cavity_data
            best_key_info = best_scheme.get("key_info", {}) if best_scheme else best_key_info

            # 8. 保存模具型腔数据（包含方案级预览链接）
            await self.storage_service.save_mold_cavity_data(
                db_session, stp_file_id, detailed_cavity_json
            )

            html_file_path = self.html_generator.generate_and_save_visualization(
                geometry_data,
                Path(file_path).name,
                cavity_data=best_cavity_data,
                pointcloud_data=pointcloud_data,
                lod_data=lod_data,
            )

            await self.storage_service.save_html_file(
                db_session,
                stp_file_id,
                Path(html_file_path).name,
                html_file_path,
            )
            stage_timings["persist_artifacts"] = round(time.perf_counter() - stage_started, 3)

            # 9. 分析模具设计
            stage_started = time.perf_counter()
            analysis_result = self.geometry_analyzer.analyze_mold_design(
                geometry_data,
                product_material=requested_material,
                shape=shape,
            )

            if analysis_result:
                await self.storage_service.save_features_and_recommendations(
                    db_session,
                    stp_file_id,
                    analysis_result.get("detected_features", []),
                    analysis_result.get("design_recommendations", []),
                )

                await self._save_analysis_metrics(db_session, stp_file_id, analysis_result)
            stage_timings["analyze_design"] = round(time.perf_counter() - stage_started, 3)

            # 9.6 更新STP文件的分析摘要字段
            await self.storage_service.update_stp_file_analysis_summary(
                db_session,
                stp_file_id,
                volume=geometry_data.get("volume", 0),
                surface_area=geometry_data.get("surface_area", 0),
                product_weight=CalculationService.calculate_product_weight(
                    geometry_data.get("volume", 0), selected_material["density"]
                ),
            )

            # 9.7 FreeCAD 几何验证
            stage_started = time.perf_counter()
            verification_result = await self._step_verify(
                file_path, db_session, task_id, stp_file_id, analysis_result
            )
            stage_timings["verify_geometry"] = round(time.perf_counter() - stage_started, 3)

            # 9.8 LLM 增强分析
            llm_report = None
            stage_started = time.perf_counter()
            if analysis_result:
                side_action_ai = await llm_service.generate_side_action_analysis(
                    analysis_result, detailed_cavity_json
                )
                design_report = await llm_service.generate_design_report(
                    analysis_result, detailed_cavity_json
                )
                llm_report = llm_service.compose_llm_report(
                    design_report, side_action_ai
                )
            stage_timings["generate_llm_report"] = round(time.perf_counter() - stage_started, 3)

            # 10. 完成处理
            await self.storage_service.update_stp_file_status(db_session, stp_file_id, "completed")
            await self.storage_service.update_task_status(
                db_session, task_id, "completed", 100, "模具型腔生成完成"
            )
            await self.storage_service.update_task_parameters(
                db_session,
                task_id,
                {
                    "stage_timings": stage_timings,
                    "material": requested_material,
                    "verification": verification_result,
                    "llm_report": llm_report,
                    "export_artifacts": export_artifacts,
                    **process_params,
                },
            )

            # 更新任务缓存状态
            await redis_task_manager.update_task(task_id, {
                "geometry_data": geometry_data,
                "analysis_result": analysis_result,
                "plan_result": detailed_cavity_json,
                "candidate_schemes": detailed_cavity_json.get("candidate_schemes", []),
                "best_scheme_id": detailed_cavity_json.get("best_scheme_id"),
                "cavity_data": best_cavity_data,
                "key_info": best_key_info,
                "material": requested_material,
                "parameters": process_params,
                "stage_timings": stage_timings,
                "html_file": best_scheme.get("html_file", f"/html/{Path(html_file_path).name}") if best_scheme else f"/html/{Path(html_file_path).name}",
                "verification": verification_result,
                "llm_report": llm_report,
                "export_artifacts": export_artifacts,
                "status": ProcessingStatus.COMPLETED,
                "completed_at": str(datetime.now()),
            })

            logger.info(f"模具型腔生成完成: {task_id}")
            logger.info(f"key_info metadata: {detailed_cavity_json.get('metadata', {})}")
            logger.info(f"key_info manufacturing_info: {detailed_cavity_json.get('manufacturing_info', {})}")
            logger.info(f"key_info geometric_characteristics: {detailed_cavity_json.get('mold_cavities', {}).get('cavity_key_info', {}).get('geometric_characteristics', {})}")

        except Exception as e:
            logger.error(f"模具型腔生成失败: {e}")

            await self.storage_service.update_stp_file_status(db_session, stp_file_id, "failed")
            await self.storage_service.update_task_status(
                db_session, task_id, "failed", error_message=str(e)
            )

            task = await redis_task_manager.get_task(task_id)
            if task:
                await redis_task_manager.update_task(task_id, {
                    "status": ProcessingStatus.FAILED,
                    "error": str(e),
                    "completed_at": str(datetime.now()),
                })

    # ─── 内部步骤 ───

    async def _step_generate_mesh(
        self, shape, geometry_data: dict, file_path: str,
        db_session: AsyncSession, stp_file_id: int, task_id: str,
    ) -> Optional[Dict[str, Any]]:
        """生成网格数据并持久化，失败不影响主流程"""
        mesh_result = None
        try:
            mesh_result = self.mesh_generator.generate_mesh_from_shape(shape)

            vertices = mesh_result.get("vertices", [])
            faces = mesh_result.get("faces", [])
            points = mesh_result.get("points", [])
            normals = mesh_result.get("normals", [])
            point_count = mesh_result.get("point_count", 0)
            vertex_count = mesh_result.get("vertex_count", 0)
            face_count = mesh_result.get("face_count", 0)

            if vertices and faces:
                bbox = geometry_data.get("bounding_box", {})

                mesh_json = {
                    "metadata": {
                        "file_name": Path(file_path).name,
                        "generated_at": datetime.now().isoformat(),
                        "quality": "medium",
                        "vertex_count": vertex_count,
                        "face_count": face_count,
                        "point_count": point_count,
                    },
                    "mesh": {
                        "vertices": vertices,
                        "faces": faces,
                    },
                    "pointcloud": {
                        "points": points,
                        "normals": normals,
                        "count": point_count,
                    },
                    "bounding_box": bbox,
                }

                await self.storage_service.save_mesh_data(
                    db_session,
                    stp_file_id=stp_file_id,
                    mesh_json=mesh_json,
                    quality="medium",
                )
                await redis_task_manager.update_task(task_id, {
                    "mesh_summary": {
                        "vertex_count": vertex_count,
                        "face_count": face_count,
                        "point_count": point_count,
                        "quality": "medium",
                    }
                })
        except Exception as mesh_err:
            logger.warning(f"网格生成或保存失败，不影响主流程: {mesh_err}")

        return mesh_result

    async def _step_generate_cavity(
        self, shape, selected_material: dict, is_foam_material: bool, process_params: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """生成多方案分模结果"""
        plan_result = None
        try:
            if shape:
                plan_result = self.multi_scheme_planner.generate_plan(
                    shape=shape,
                    material=selected_material,
                    is_foam_material=is_foam_material,
                    process_params=process_params,
                )
                logger.info(
                    f"多方案分模完成: 生成 {len(plan_result.get('candidate_schemes', []))} 套方案"
                )
        except Exception as cavity_err:
            logger.warning(f"多方案分模失败，使用简化数据: {cavity_err}")
            traceback.print_exc()
            plan_result = None

        return plan_result

    def _cache_export_shapes(self, task_id: str, export_shapes: Dict[str, Dict[str, Any]]):
        self._export_shapes_cache[task_id] = export_shapes

    def _persist_step_exports(
        self,
        task_id: str,
        original_filename: str,
        export_shapes: Dict[str, Dict[str, Any]],
    ) -> Optional[Dict[str, Any]]:
        if not export_shapes:
            return None

        base_filename = Path(original_filename).stem or f"mold_{task_id}"
        manifest = {
            "version": 1,
            "task_id": task_id,
            "generated_at": datetime.now().isoformat(),
            "schemes": {},
        }
        components = ["cavity", "core", "parting_surface"]

        for scheme_id, cavity_data in export_shapes.items():
            try:
                result = self.cad_exporter.export_mold_results(
                    cavity_data=cavity_data,
                    base_filename=base_filename,
                    formats=["step"],
                    components=components,
                    task_id=task_id,
                    scheme_id=scheme_id,
                )
                manifest["schemes"][scheme_id] = {
                    "base_filename": result.get("base_filename"),
                    "generated_at": datetime.now().isoformat(),
                    "files": result.get("files", []),
                    "errors": result.get("errors", []),
                    "total_files": result.get("total_files", 0),
                    "total_errors": result.get("total_errors", 0),
                }
            except Exception as exc:
                logger.warning("持久化 STEP 导出失败: task=%s scheme=%s error=%s", task_id, scheme_id, exc)
                manifest["schemes"][scheme_id] = {
                    "base_filename": base_filename,
                    "generated_at": datetime.now().isoformat(),
                    "files": [],
                    "errors": [str(exc)],
                    "total_files": 0,
                    "total_errors": 1,
                }

        return manifest

    def get_export_shapes(self, task_id: str, scheme_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
        scheme_map = self._export_shapes_cache.get(task_id, {})
        if not scheme_map:
            return None
        if scheme_id:
            return scheme_map.get(scheme_id)
        return next(iter(scheme_map.values()), None)

    @staticmethod
    def _normalize_process_params(process_params: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        payload = dict(process_params or {})
        return {
            "material": MaterialService.resolve_material(str(payload.get("material", "ABS"))),
            "draft_angle": float(payload.get("draft_angle", 2.0)),
            "shrinkage_rate": float(payload.get("shrinkage_rate", 0.5)),
            "parting_precision": float(payload.get("parting_precision", 0.1)),
            "cavity_match": int(payload.get("cavity_match", 95)),
        }

    async def _step_verify(
        self, file_path: str, db_session: AsyncSession,
        task_id: str, stp_file_id: int, analysis_result: Optional[dict],
    ) -> Optional[Dict[str, Any]]:
        """FreeCAD 几何验证（可通过配置禁用）"""
        from config.settings import settings

        if not settings.ENABLE_FREECAD_VERIFICATION:
            logger.info("FreeCAD验证已禁用（设置 ENABLE_FREECAD_VERIFICATION=true 启用）")
            return {"status": "disabled", "reason": "FreeCAD验证已禁用"}

        await self.storage_service.update_task_status(
            db_session, task_id, "processing", 90, "FreeCAD几何验证"
        )

        try:
            from services.verification_service import GeometryVerificationService
            verification_svc = GeometryVerificationService(timeout=settings.FREECAD_VERIFICATION_TIMEOUT)
            verification_result = await verification_svc.verify_stp_file(file_path)

            if verification_result and analysis_result:
                await self._save_verification_metrics(db_session, stp_file_id, verification_result)

            logger.info(f"FreeCAD验证完成: {verification_result.get('status', 'unknown') if verification_result else 'failed'}")
            return verification_result
        except Exception as ve:
            logger.warning(f"FreeCAD验证失败（不影响主流程）: {ve}")
            return {"status": "error", "error": str(ve)}

    async def _attach_scheme_previews(
        self,
        detailed_cavity_json: Dict[str, Any],
        geometry_data: Dict[str, Any],
        stp_filename: str,
        pointcloud_data: Optional[Dict[str, Any]] = None,
        lod_data: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """为每个候选分模方案生成独立HTML预览链接。"""
        candidate_schemes = detailed_cavity_json.get("candidate_schemes", [])
        if not candidate_schemes:
            return detailed_cavity_json

        for scheme in candidate_schemes:
            cavity_data = scheme.get("cavity_data")
            if not cavity_data:
                continue
            suffix = scheme.get("scheme_id")
            html_path = self.html_generator.generate_and_save_visualization(
                geometry_data,
                stp_filename,
                cavity_data=cavity_data,
                pointcloud_data=pointcloud_data,
                suffix=suffix,
                lod_data=lod_data,
            )
            scheme["html_file"] = f"/html/{Path(html_path).name}"

        best_scheme = CalculationService.get_best_scheme(detailed_cavity_json)
        if best_scheme:
            detailed_cavity_json["html_file"] = best_scheme.get("html_file")

        return detailed_cavity_json

    # ─── 指标持久化 ───

    async def _save_analysis_metrics(self, session: AsyncSession, stp_file_id: int, analysis_result: dict):
        """保存分析指标到数据库"""
        from models.database import AnalysisMetrics

        quality_metrics = analysis_result.get("quality_metrics", {})
        analysis_summary = analysis_result.get("analysis_summary", "")

        metrics = AnalysisMetrics(
            stp_file_id=stp_file_id,
            volume_utilization=quality_metrics.get("volume_utilization", 0),
            topology_complexity=quality_metrics.get("topology_complexity", 0),
            wall_uniformity=quality_metrics.get("wall_uniformity", 0),
            analysis_summary=analysis_summary,
        )

        session.add(metrics)
        await session.commit()
        logger.info(f"分析指标保存成功: {metrics.id}")

    async def _save_verification_metrics(self, session: AsyncSession, stp_file_id: int, verification_result: dict):
        """保存验证指标到数据库"""
        from models.database import AnalysisMetrics
        from sqlalchemy import select

        result = await session.execute(
            select(AnalysisMetrics).where(AnalysisMetrics.stp_file_id == stp_file_id)
        )
        metrics = result.scalar_one_or_none()

        comparison = verification_result.get("comparison", {})
        volume_comparison = comparison.get("volume", {})
        area_comparison = comparison.get("surface_area", {})

        if metrics:
            metrics.verification_status = verification_result.get("status", "unknown")
            metrics.verification_volume_diff = volume_comparison.get("difference_percent", 0)
            metrics.verification_area_diff = area_comparison.get("difference_percent", 0)
            metrics.verification_details = verification_result
        else:
            metrics = AnalysisMetrics(
                stp_file_id=stp_file_id,
                verification_status=verification_result.get("status", "unknown"),
                verification_volume_diff=volume_comparison.get("difference_percent", 0),
                verification_area_diff=area_comparison.get("difference_percent", 0),
                verification_details=verification_result,
            )
            session.add(metrics)

        await session.commit()
        logger.info(f"验证指标保存成功: stp_file_id={stp_file_id}")


# 模块级单例，供路由层直接使用
processing_service = ProcessingService()