geMoldInsight/src/moldinsight/services/analysis_storage_service.py

# services/analysis_storage_service.py
"""分析结果数据存储——几何/网格/型腔/HTML/特征的 RustFS 上传与 PG 元数据，
以及任务完整数据视图的组装。

批次 3 自 storage_integration_rustfs.py 按职责拆分（见 task_storage_service.py 头注）。
"""
import json
from pathlib import Path
from typing import Optional, Dict, Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import joinedload

from moldinsight.models import STPFile, GeometryData, MeshData, MoldCavityData, HTMLFile, FeatureDetection, DesignRecommendation
from moldinsight.storage.rustfs_storage import rustfs_manager
from shared.utils.logger import get_logger

logger = get_logger(__name__)


class AnalysisStorageService:
    """分析结果数据（几何/网格/型腔/HTML/特征）存储与视图组装"""

    @staticmethod
    def _resolve_best_scheme_payload(cavity_json: Dict[str, Any]) -> Dict[str, Any]:
        """从新多方案/旧单方案结构中解析推荐方案和型腔详情。"""
        if not isinstance(cavity_json, dict):
            return {
                "best_scheme_id": None,
                "best_scheme": {},
                "best_cavity_data": {},
                "key_info": {},
            }

        candidate_schemes = cavity_json.get("candidate_schemes") or []
        if not candidate_schemes:
            key_info = cavity_json.get("mold_cavities", {}).get("cavity_key_info", {})
            return {
                "best_scheme_id": cavity_json.get("best_scheme_id"),
                "best_scheme": {},
                "best_cavity_data": cavity_json,
                "key_info": key_info,
            }

        best_scheme_id = cavity_json.get("best_scheme_id")
        best_scheme = candidate_schemes[0]
        if best_scheme_id:
            for scheme in candidate_schemes:
                if scheme.get("scheme_id") == best_scheme_id:
                    best_scheme = scheme
                    break

        best_cavity_data = best_scheme.get("cavity_data", {}) if isinstance(best_scheme, dict) else {}
        key_info = best_scheme.get("key_info", {}) if isinstance(best_scheme, dict) else {}
        if not key_info:
            key_info = best_cavity_data.get("mold_cavities", {}).get("cavity_key_info", {})
        return {
            "best_scheme_id": best_scheme.get("scheme_id") or best_scheme_id,
            "best_scheme": best_scheme,
            "best_cavity_data": best_cavity_data,
            "key_info": key_info,
        }

    @staticmethod
    def _parse_first_number(value: Any) -> Optional[float]:
        if value is None:
            return None
        try:
            return float(value)
        except (TypeError, ValueError):
            pass
        import re
        matches = re.findall(r"\d+(?:\.\d+)?", str(value))
        if not matches:
            return None
        try:
            return float(matches[0])
        except (TypeError, ValueError):
            return None

    async def save_geometry_data(self, session: AsyncSession,
                               stp_file_id: int,
                               geometry_json: Dict[str, Any],
                               analysis_method: str = "pythonocc") -> GeometryData:
        """保存几何数据到PostgreSQL元数据 + RustFS对象存储"""

        # 1. 获取文件哈希
        stp_file = await session.get(STPFile, stp_file_id)
        file_hash = stp_file.file_hash

        # 2. 上传到RustFS
        upload_result = await rustfs_manager.upload_json_data(
            file_type='geometry_data',
            json_data=geometry_json,
            file_hash=file_hash
        )

        # 3. 提取几何数据
        if 'geometry_data' in geometry_json:
            geo_data = geometry_json['geometry_data']
        else:
            geo_data = geometry_json

        # 4. 创建PostgreSQL记录
        geometry_data = GeometryData(
            stp_file_id=stp_file_id,
            object_key=upload_result['object_key'],
            storage_bucket=upload_result['bucket'],
            analysis_method=analysis_method,

            # 提取摘要字段
            volume=geo_data.get('volume'),
            surface_area=geo_data.get('surface_area'),
            bounding_box_min=geo_data.get('bounding_box', {}).get('min'),
            bounding_box_max=geo_data.get('bounding_box', {}).get('max'),
            center_of_mass=geo_data.get('center_of_mass'),
            topology_faces=geo_data.get('topology', {}).get('faces'),
            topology_edges=geo_data.get('topology', {}).get('edges'),
            topology_vertices=geo_data.get('topology', {}).get('vertices')
        )

        session.add(geometry_data)
        # D9：数据本体仅 flush，与网格等同阶段数据由编排层统一 commit（原子落库）
        await session.flush()
        await session.refresh(geometry_data)

        logger.info(f"几何数据保存成功 RustFS: {geometry_data.id}")
        return geometry_data

    async def save_mesh_data(
        self,
        session: AsyncSession,
        stp_file_id: int,
        mesh_json: Dict[str, Any],
        quality: str = "medium"
    ) -> MeshData:
        """保存网格数据到 PostgreSQL 元数据 + RustFS 对象存储

        mesh_json 为完整网格 JSON（顶点、面、点云等），
        PostgreSQL 只存 object_key 和一些摘要字段，详细数据放在 RustFS。
        """

        # 1. 获取文件哈希
        stp_file = await session.get(STPFile, stp_file_id)
        file_hash = stp_file.file_hash

        # 2. 上传网格 JSON 到 RustFS
        upload_result = await rustfs_manager.upload_json_data(
            file_type='mesh_data',
            json_data=mesh_json,
            file_hash=file_hash
        )

        # 3. 提取摘要信息
        mesh_section = mesh_json.get('mesh', {})
        pointcloud_section = mesh_json.get('pointcloud', {})
        bbox = mesh_json.get('bounding_box', {})

        vertices = mesh_section.get('vertices') or []
        faces = mesh_section.get('faces') or []

        vertex_count = len(vertices)
        face_count = len(faces)
        point_count = pointcloud_section.get('count')

        # 4. 创建 PostgreSQL 记录
        mesh_data = MeshData(
            stp_file_id=stp_file_id,
            object_key=upload_result['object_key'],
            storage_bucket=upload_result['bucket'],
            quality=quality,
            vertex_count=vertex_count,
            face_count=face_count,
            point_count=point_count,
            bounding_box_min=bbox.get('min'),
            bounding_box_max=bbox.get('max')
        )

        session.add(mesh_data)
        # D9：数据本体仅 flush，与几何数据同阶段由编排层统一 commit
        await session.flush()
        await session.refresh(mesh_data)

        logger.info(f"网格数据保存成功 RustFS: {mesh_data.id}")
        return mesh_data

    async def save_mold_cavity_data(self, session: AsyncSession,
                                    stp_file_id: int,
                                    cavity_json: Dict[str, Any]) -> MoldCavityData:
        """保存模具型腔数据到PostgreSQL元数据 + RustFS对象存储"""

        # 1. 获取文件哈希
        stp_file = await session.get(STPFile, stp_file_id)
        file_hash = stp_file.file_hash

        # 2. 上传到RustFS
        upload_result = await rustfs_manager.upload_json_data(
            file_type='mold_cavities',
            json_data=cavity_json,
            file_hash=file_hash
        )

        # 3. 提取关键信息（兼容多方案与单方案结构）
        payload = self._resolve_best_scheme_payload(cavity_json)
        best_scheme_id = payload.get("best_scheme_id")
        best_scheme = payload.get("best_scheme") or {}
        best_cavity_data = payload.get("best_cavity_data") or {}
        metadata = best_cavity_data.get('metadata', {})
        product_analysis = best_cavity_data.get('product_analysis', {})
        manufacturing_info = best_cavity_data.get('manufacturing_info', {})
        mold_size = manufacturing_info.get('estimated_mold_size', {})
        key_info = payload.get("key_info") or {}
        if not key_info:
            key_info = best_cavity_data.get('mold_cavities', {}).get('cavity_key_info', {})
        mold_material = (
            metadata.get("selected_material")
            or manufacturing_info.get("recommended_material")
            or 'Aluminum Alloy 7075'
        )
        parting_line_length = self._parse_first_number(
            manufacturing_info.get("parting_line_length")
        )

        # 4. 创建PostgreSQL记录
        mold_cavity = MoldCavityData(
            stp_file_id=stp_file_id,
            detailed_object_key=upload_result['object_key'],
            storage_bucket=upload_result['bucket'],

            # 模具参数
            mold_material=mold_material,
            shrinkage_rate=metadata.get('shrinkage_rate', 0.005),
            draft_angle=metadata.get('draft_angle', 2.0),
            parting_line_length=parting_line_length,

            # 提取的摘要字段
            cavity_key_info=key_info,
            mold_size_length=mold_size.get('length'),
            mold_size_width=mold_size.get('width'),
            mold_size_height=mold_size.get('height'),
            estimated_clamping_force=manufacturing_info.get('estimated_clamping_force'),
            product_volume=product_analysis.get('volume'),

            # 从key_info中提取（如果存在）
            product_weight=key_info.get('geometric_characteristics', {}).get('product_weight'),
            wall_thickness_range=key_info.get('geometric_characteristics', {}).get('wall_thickness_range'),
            complexity_score=key_info.get('geometric_characteristics', {}).get('complexity_score'),

            # 质量评估
            weld_line_risk=key_info.get('quality_considerations', {}).get('potential_weld_lines'),
            sink_mark_risk=key_info.get('quality_considerations', {}).get('sink_mark_areas'),
            warpage_risk=key_info.get('quality_considerations', {}).get('warpage_risk'),

            # 多方案可信化摘要
            best_scheme_id=best_scheme_id,
            confidence_score=best_scheme.get("confidence_score"),
            is_fallback=best_scheme.get("is_fallback"),
            fallback_reason=best_scheme.get("fallback_reason"),
        )

        session.add(mold_cavity)
        # D9：数据本体仅 flush，型腔/HTML/特征同属结果包，由编排层统一 commit
        await session.flush()
        await session.refresh(mold_cavity)

        logger.info(f"模具型腔数据保存成功 RustFS: {mold_cavity.id}")
        return mold_cavity

    async def save_html_file(self, session: AsyncSession,
                            stp_file_id: int,
                            filename: str,
                            file_path: str,
                            visualization_type: str = "3d_viewer") -> HTMLFile:
        """保存HTML文件：正文按报告键裸传 RustFS + PG 仅存元数据（D11）。

        上传键固定 html/reports/{filename}（文件名寻址），读侧 GET /html/{filename}
        按文件名直取，不再写 JSON 包装对象（遗留 html/{hash}.json 仅由读侧兼容解析）。
        file_path 为任务内临时目录路径，任务结束即删除，仅供追溯，不作为读取来源。
        """
        html_file_path = Path(file_path)
        if not html_file_path.exists():
            raise RuntimeError(f"HTML 可视化文件缺失: {file_path}")
        html_bytes = html_file_path.read_bytes()

        object_key = await rustfs_manager.upload_report_artifact(
            filename, html_bytes, content_type="text/html; charset=utf-8"
        )

        html_file = HTMLFile(
            stp_file_id=stp_file_id,
            object_key=object_key,
            storage_bucket=rustfs_manager.bucket_name,
            filename=filename,
            file_path=file_path,
            # 停止双写完整 HTML 进 PG：读取路径走 /html/{filename} 代理，
            # PG 仅存对象键与文件名，避免大文本撑爆表
            html_content=None,
            visualization_type=visualization_type
        )

        session.add(html_file)
        # D9：数据本体仅 flush，型腔/HTML/特征同属结果包，由编排层统一 commit
        await session.flush()
        await session.refresh(html_file)

        logger.info(f"HTML文件保存成功 RustFS: {html_file.id} ({object_key})")
        return html_file

    async def save_features_and_recommendations(
        self, session: AsyncSession,
        stp_file_id: int,
        features: list,
        recommendations: list
    ):
        """保存特征检测结果和设计建议"""

        # 1. 保存特征
        for feature in features:
            feature_record = FeatureDetection(
                stp_file_id=stp_file_id,
                feature_type=feature.get('feature_type'),
                confidence=feature.get('confidence'),
                location=feature.get('location'),
                dimensions=feature.get('dimensions'),
                parameters=feature.get('parameters')
            )
            session.add(feature_record)

        # 2. 保存建议
        for rec in recommendations:
            rec_record = DesignRecommendation(
                stp_file_id=stp_file_id,
                rec_type=rec.get('type') or rec.get('rec_type'),
                priority=rec.get('priority'),
                description=rec.get('description'),
                reason=rec.get('reason'),
                parameters=rec.get('parameters')
            )
            session.add(rec_record)

        # D9：数据本体仅 flush，型腔/HTML/特征同属结果包，由编排层统一 commit
        await session.flush()
        logger.info(f"保存了 {len(features)} 个特征和 {len(recommendations)} 个建议")

    async def get_stp_file_with_data(self, session: AsyncSession,
                                    stp_file_id: int) -> Dict[str, Any]:
        """获取STP文件及其所有关联数据"""
        try:
            # 1. 获取STP文件记录（使用 joinedload 预加载关联数据）
            result = await session.execute(
                select(STPFile).options(
                    joinedload(STPFile.geometry_data),
                    joinedload(STPFile.mesh_data),
                    joinedload(STPFile.mold_cavity_data),
                    joinedload(STPFile.html_file),
                    joinedload(STPFile.analysis_metrics)
                ).where(STPFile.id == stp_file_id)
            )
            stp_file = result.scalar_one_or_none()
            if not stp_file:
                raise ValueError(f"STP文件不存在: {stp_file_id}")
        except Exception as e:
            logger.error(f"获取STP文件记录失败: {e}")
            raise

        result = {
            'metadata': {
                'id': stp_file.id,
                'original_filename': stp_file.original_filename,
                'file_size': stp_file.file_size,
                'file_hash': stp_file.file_hash,
                'upload_time': stp_file.upload_time.isoformat() if stp_file.upload_time else None,
                'status': stp_file.status,
                'user_id': stp_file.user_id
            },
            'geometry_data': None,
            'mesh_data': None,
            'mold_cavity_data': None,
            'features': [],
            'recommendations': [],
            'analysis_metrics': None  # 新增分析指标字段
        }

        # 2. 从RustFS获取数据
        try:
            # 几何数据
            if stp_file.geometry_data:
                geo_data_bytes = await rustfs_manager.download_file(
                    file_type='geometry_data',
                    object_key=stp_file.geometry_data.object_key
                )
                result['geometry_data'] = json.loads(geo_data_bytes.decode('utf-8'))

            # 模具型腔数据
            if stp_file.mold_cavity_data:
                cavity_data_bytes = await rustfs_manager.download_file(
                    file_type='mold_cavities',
                    object_key=stp_file.mold_cavity_data.detailed_object_key
                )
                result['mold_cavity_data'] = json.loads(cavity_data_bytes.decode('utf-8'))

            # 网格数据
            if stp_file.mesh_data:
                mesh_bytes = await rustfs_manager.download_file(
                    file_type='mesh_data',
                    object_key=stp_file.mesh_data.object_key
                )
                result['mesh_data'] = json.loads(mesh_bytes.decode('utf-8'))

            # HTML 正文不再随本视图下载（D11）：历史实现把整个 HTML 读进
            # result['html_content']，但所有调用方只取 geometry/cavity/features，
            # HTML 的读取入口是 /html/{filename} 代理路由——每次任务查询白下载
            # 数 MB 正文纯属浪费，已删除
        except Exception as e:
            logger.error(f"从RustFS获取数据失败: {e}")

        # 3. 从PostgreSQL获取特征和建议
        features = await session.execute(
            select(FeatureDetection).where(FeatureDetection.stp_file_id == stp_file_id)
        )
        result['features'] = [
            {
                'feature_type': f.feature_type,
                'confidence': f.confidence,
                'location': f.location,
                'dimensions': f.dimensions,
                'parameters': f.parameters
            }
            for f in features.scalars().all()
        ]

        recommendations = await session.execute(
            select(DesignRecommendation).where(DesignRecommendation.stp_file_id == stp_file_id)
        )
        result['recommendations'] = [
            {
                'type': r.rec_type,  # 改为 type 以匹配前端期望的字段名
                'priority': r.priority,
                'description': r.description,
                'reason': r.reason,
                'parameters': r.parameters
            }
            for r in recommendations.scalars().all()
        ]

        # 4. 获取分析指标
        if stp_file.analysis_metrics:
            result['analysis_metrics'] = {
                'volume_utilization': stp_file.analysis_metrics.volume_utilization,
                'topology_complexity': stp_file.analysis_metrics.topology_complexity,
                'wall_uniformity': stp_file.analysis_metrics.wall_uniformity,
                'analysis_summary': stp_file.analysis_metrics.analysis_summary,
                'verification_status': stp_file.analysis_metrics.verification_status,
                'verification_volume_diff': stp_file.analysis_metrics.verification_volume_diff,
                'verification_area_diff': stp_file.analysis_metrics.verification_area_diff,
                'verification_details': stp_file.analysis_metrics.verification_details,
            }

        return result

    async def delete_stp_file_cascade(self, session: AsyncSession,
                                     stp_file_id: int):
        """级联删除STP文件及其所有关联数据"""

        stp_file = await session.get(STPFile, stp_file_id)
        if not stp_file:
            raise ValueError(f"STP文件不存在: {stp_file_id}")

        # 1. 删除RustFS中的文件
        try:
            if stp_file.object_key:
                await rustfs_manager.delete_file('stp_files', stp_file.object_key)
        except Exception as e:
            logger.error(f"删除RustFS文件失败: {e}")

        try:
            if stp_file.geometry_data:
                await rustfs_manager.delete_file('geometry_data', stp_file.geometry_data.object_key)
        except Exception as e:
            logger.error(f"删除几何数据失败: {e}")

        try:
            if stp_file.mold_cavity_data:
                await rustfs_manager.delete_file('mold_cavities', stp_file.mold_cavity_data.detailed_object_key)
        except Exception as e:
            logger.error(f"删除型腔数据失败: {e}")

        try:
            if stp_file.mesh_data:
                await rustfs_manager.delete_file('mesh_data', stp_file.mesh_data.object_key)
        except Exception as e:
            logger.error(f"删除网格数据失败: {e}")

        try:
            if stp_file.html_file:
                await rustfs_manager.delete_file('html_files', stp_file.html_file.object_key)
        except Exception as e:
            logger.error(f"删除HTML文件失败: {e}")

        # 2. 级联删除PostgreSQL记录（通过外键自动处理）
        await session.delete(stp_file)
        await session.commit()

        logger.info(f"STP文件及其关联数据已删除: {stp_file_id}")