geMoldInsight/src/core/cad_exporter.py

"""
CAD 文件导出模块

支持导出格式：
1. STEP (ISO 10303) - 推荐，UG/NX、FreeCAD、SolidWorks 通用
2. IGES (Initial Graphics Exchange Specification) - 兼容旧系统
3. STL (STereoLithography) - 网格格式，3D打印/快速预览
4. BRep (Boundary Representation) - OpenCASCADE 原生格式

导出内容：
- 型腔 (Cavity)
- 型芯 (Core)
- 分型面 (Parting Surface)
- A板/B板
- 模具块
- 完整模具装配体（多形状合并）

UG/NX 导入建议：
- 优先使用 STEP AP214 或 AP242 格式
- IGES 作为备选
- STL 仅用于预览，不可编辑

FreeCAD 导入建议：
- STEP AP214 最佳兼容性
- BRep 可直接在 FreeCAD 的 OpenCASCADE 内核中打开
"""

import os
import re
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path
from OCC.Core.TopoDS import TopoDS_Shape
from utils.logger import get_logger

logger = get_logger(__name__)


class CADExporter:
    """CAD 文件导出器"""

    def __init__(self, output_dir: str = "./exports"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)

    @staticmethod
    def _safe_segment(value: Optional[str], fallback: str) -> str:
        text = str(value or "").strip()
        if not text:
            text = fallback
        text = re.sub(r"[^A-Za-z0-9._-]+", "_", text)
        return text[:80] or fallback

    def build_export_dir(
        self,
        base_filename: str,
        task_id: Optional[str] = None,
        scheme_id: Optional[str] = None,
    ) -> str:
        if task_id:
            task_segment = self._safe_segment(task_id, "task")
            scheme_segment = self._safe_segment(scheme_id, "default")
            return os.path.join(self.output_dir, task_segment, scheme_segment)
        return os.path.join(self.output_dir, self._safe_segment(base_filename, "mold"))

    def get_relative_path(self, filepath: str) -> str:
        full_path = Path(filepath).resolve()
        output_root = Path(self.output_dir).resolve()
        try:
            relative = full_path.relative_to(output_root)
        except ValueError:
            relative = Path(os.path.basename(filepath))
        return relative.as_posix()

    def export_step(self, shape: TopoDS_Shape, filepath: str,
                     schema: str = "AP214") -> bool:
        """
        导出 STEP 文件

        Args:
            shape: OpenCASCADE TopoDS_Shape
            filepath: 输出文件路径
            schema: STEP 应用协议 (AP203/AP214/AP242)

        Returns:
            是否成功
        """
        try:
            from OCC.Core.STEPControl import (
                STEPControl_Writer,
                STEPControl_AsIs,
            )
            from OCC.Core.Interface import Interface_Static

            writer = STEPControl_Writer()

            if schema == "AP203":
                Interface_Static.SetCVal("write.step.schema", "AP203")
            elif schema == "AP242":
                Interface_Static.SetCVal("write.step.schema", "AP242")
            else:
                Interface_Static.SetCVal("write.step.schema", "AP214")

            writer.Transfer(shape, STEPControl_AsIs)

            status = writer.Write(filepath)

            if status == 1:
                file_size = os.path.getsize(filepath) if os.path.exists(filepath) else 0
                logger.info(f"STEP 导出成功: {filepath} ({file_size} bytes, {schema})")
                return True
            else:
                logger.error(f"STEP 导出失败: 写入状态={status}")
                return False

        except ImportError as e:
            logger.error(f"STEP 导出依赖缺失: {e}")
            return False
        except Exception as e:
            logger.error(f"STEP 导出失败: {e}")
            return False

    def export_iges(self, shape: TopoDS_Shape, filepath: str) -> bool:
        """
        导出 IGES 文件

        Args:
            shape: OpenCASCADE TopoDS_Shape
            filepath: 输出文件路径

        Returns:
            是否成功
        """
        try:
            from OCC.Core.IGESControl import IGESControl_Writer
            from OCC.Core.Interface import Interface_Static

            Interface_Static.SetCVal("write.iges.brep.mode", "0")

            writer = IGESControl_Writer()
            writer.AddShape(shape)
            writer.ComputeModel()

            status = writer.Write(filepath)

            if status:
                file_size = os.path.getsize(filepath) if os.path.exists(filepath) else 0
                logger.info(f"IGES 导出成功: {filepath} ({file_size} bytes)")
                return True
            else:
                logger.error("IGES 导出失败: 写入返回 False")
                return False

        except ImportError as e:
            logger.error(f"IGES 导出依赖缺失: {e}")
            return False
        except Exception as e:
            logger.error(f"IGES 导出失败: {e}")
            return False

    def export_stl(self, shape: TopoDS_Shape, filepath: str,
                    ascii_mode: bool = True,
                    deflection: float = 0.1) -> bool:
        """
        导出 STL 文件

        Args:
            shape: OpenCASCADE TopoDS_Shape
            filepath: 输出文件路径
            ascii_mode: True=ASCII格式, False=二进制格式
            deflection: 网格偏差（越小越精细）

        Returns:
            是否成功
        """
        try:
            from OCC.Core.StlAPI import StlAPI_Writer
            from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh

            mesh = BRepMesh_IncrementalMesh(shape, deflection)
            mesh.Perform()

            if not mesh.IsDone():
                logger.warning("STL 网格化未完成，尝试继续导出")

            writer = StlAPI_Writer()
            writer.AsciiMode = ascii_mode

            writer.Write(shape, filepath)

            if os.path.exists(filepath) and os.path.getsize(filepath) > 0:
                file_size = os.path.getsize(filepath)
                logger.info(f"STL 导出成功: {filepath} ({file_size} bytes)")
                return True
            else:
                logger.error("STL 导出失败: 文件为空或不存在")
                return False

        except ImportError as e:
            logger.error(f"STL 导出依赖缺失: {e}")
            return False
        except Exception as e:
            logger.error(f"STL 导出失败: {e}")
            return False

    def export_brep(self, shape: TopoDS_Shape, filepath: str) -> bool:
        """
        导出 BRep 文件（OpenCASCADE 原生格式）

        FreeCAD 可直接导入此格式

        Args:
            shape: OpenCASCADE TopoDS_Shape
            filepath: 输出文件路径

        Returns:
            是否成功
        """
        try:
            from OCC.Core.BRepTools import BRepTools_Write

            BRepTools_Write(shape, filepath)

            if os.path.exists(filepath) and os.path.getsize(filepath) > 0:
                file_size = os.path.getsize(filepath)
                logger.info(f"BRep 导出成功: {filepath} ({file_size} bytes)")
                return True
            else:
                logger.error("BRep 导出失败: 文件为空或不存在")
                return False

        except ImportError as e:
            logger.error(f"BRep 导出依赖缺失: {e}")
            return False
        except Exception as e:
            logger.error(f"BRep 导出失败: {e}")
            return False

    def export_mold_results(self, cavity_data: Dict,
                             base_filename: str,
                             formats: List[str] = None,
                             components: List[str] = None,
                             task_id: Optional[str] = None,
                             scheme_id: Optional[str] = None) -> Dict[str, Any]:
        if formats is None:
            formats = ["step", "stl"]
        if components is None:
            components = ["cavity", "core"]

        export_dir = self.build_export_dir(
            base_filename=base_filename,
            task_id=task_id,
            scheme_id=scheme_id,
        )
        os.makedirs(export_dir, exist_ok=True)

        results = {
            "base_filename": base_filename,
            "task_id": task_id,
            "scheme_id": scheme_id,
            "export_dir": export_dir,
            "files": [],
            "errors": [],
        }

        shape_map = {
            "cavity": ("cavity", "型腔"),
            "core": ("core", "型芯"),
            "parting_surface": ("parting_surface", "分型面"),
        }

        shapes_to_export: List[Tuple[str, str, TopoDS_Shape]] = []
        assembly_shapes: List[Tuple[TopoDS_Shape, str]] = []

        for comp in components:
            if comp == "all":
                for key, (data_key, label) in shape_map.items():
                    shape = cavity_data.get(data_key)
                    if shape is not None:
                        shapes_to_export.append((key, label, shape))
                        assembly_shapes.append((shape, label))
                break
            elif comp in shape_map:
                data_key, label = shape_map[comp]
                shape = cavity_data.get(data_key)
                if shape is not None:
                    shapes_to_export.append((comp, label, shape))
                    assembly_shapes.append((shape, label))
                else:
                    results["errors"].append(f"{label}形状不可用")

        # STEP: 所有组件合并为一个装配体文件
        if "step" in formats and assembly_shapes:
            filepath = os.path.join(export_dir, f"{base_filename}_mold.step")
            success = self.export_assembly_step(assembly_shapes, filepath)
            if success:
                file_size = os.path.getsize(filepath)
                relative_path = self.get_relative_path(filepath)
                results["files"].append({
                    "component": "assembly",
                    "component_label": "模具装配体",
                    "format": "step",
                    "filepath": filepath,
                    "relative_path": relative_path,
                    "filename": os.path.basename(filepath),
                    "size_bytes": file_size,
                    "size_readable": self._format_file_size(file_size),
                })
            else:
                results["errors"].append("装配体 STEP 导出失败")

        # IGES / STL / BRep: 逐组件导出
        non_assembly_formats = [f for f in formats if f != "step"]
        for comp_name, label, shape in shapes_to_export:
            for fmt in non_assembly_formats:
                filepath = os.path.join(export_dir, f"{base_filename}_{comp_name}.{fmt}")

                success = False
                if fmt == "iges":
                    success = self.export_iges(shape, filepath)
                elif fmt == "stl":
                    success = self.export_stl(shape, filepath)
                elif fmt == "brep":
                    success = self.export_brep(shape, filepath)
                else:
                    results["errors"].append(f"不支持的格式: {fmt}")
                    continue

                if success:
                    file_size = os.path.getsize(filepath)
                    relative_path = self.get_relative_path(filepath)
                    results["files"].append({
                        "component": comp_name,
                        "component_label": label,
                        "format": fmt,
                        "filepath": filepath,
                        "relative_path": relative_path,
                        "filename": os.path.basename(filepath),
                        "size_bytes": file_size,
                        "size_readable": self._format_file_size(file_size),
                    })
                else:
                    results["errors"].append(f"{label} ({fmt}) 导出失败")

        results["total_files"] = len(results["files"])
        results["total_errors"] = len(results["errors"])

        logger.info(f"模具导出完成: {results['total_files']} 个文件, "
                    f"{results['total_errors']} 个错误")

        return results

    def export_assembly_step(self, shapes_with_names: List[Tuple[TopoDS_Shape, str]],
                              filepath: str,
                              schema: str = "AP214") -> bool:
        """
        导出装配体 STEP 文件（多个形状写入同一个 STEP 文件）

        UG/NX 和 FreeCAD 可以识别装配体中的各个零件

        Args:
            shapes_with_names: [(shape, name), ...] 形状和名称列表
            filepath: 输出文件路径
            schema: STEP 协议版本

        Returns:
            是否成功
        """
        try:
            from OCC.Core.STEPControl import (
                STEPControl_Writer,
                STEPControl_AsIs,
            )
            from OCC.Core.Interface import Interface_Static

            writer = STEPControl_Writer()

            if schema == "AP203":
                Interface_Static.SetCVal("write.step.schema", "AP203")
            elif schema == "AP242":
                Interface_Static.SetCVal("write.step.schema", "AP242")
            else:
                Interface_Static.SetCVal("write.step.schema", "AP214")

            for shape, name in shapes_with_names:
                try:
                    writer.Transfer(shape, STEPControl_AsIs)
                    logger.info(f"已添加到装配体: {name}")
                except Exception as e:
                    logger.warning(f"添加形状 {name} 失败: {e}")

            status = writer.Write(filepath)

            if status == 1:
                file_size = os.path.getsize(filepath) if os.path.exists(filepath) else 0
                logger.info(f"装配体 STEP 导出成功: {filepath} ({file_size} bytes)")
                return True
            else:
                logger.error(f"装配体 STEP 导出失败: 状态={status}")
                return False

        except Exception as e:
            logger.error(f"装配体 STEP 导出失败: {e}")
            return False

    def get_export_recommendations(self, target_software: str = "ug") -> Dict[str, Any]:
        """
        获取针对目标软件的导出建议

        Args:
            target_software: 目标软件 (ug/nx, freecad, solidworks, autocad)

        Returns:
            导出建议
        """
        recommendations = {
            "ug": {
                "name": "UG/NX",
                "primary_format": "step",
                "step_schema": "AP242",
                "secondary_format": "iges",
                "notes": [
                    "推荐 STEP AP242 格式，支持颜色和装配信息",
                    "IGES 作为备选，但可能丢失拓扑信息",
                    "STL 仅用于预览，不可参数化编辑",
                    "导入时选择 '保留原始坐标系'",
                ],
                "import_settings": {
                    "step": "File → Import → STEP203/214/242",
                    "iges": "File → Import → IGES",
                    "stl": "File → Import → STL (仅可视化)",
                },
            },
            "freecad": {
                "name": "FreeCAD",
                "primary_format": "step",
                "step_schema": "AP214",
                "secondary_format": "brep",
                "notes": [
                    "STEP AP214 最佳兼容性",
                    "BRep 是 OpenCASCADE 原生格式，FreeCAD 可直接打开",
                    "导入后可在 Part 工作台中编辑",
                    "推荐使用 FreeCAD 0.21+ 版本",
                ],
                "import_settings": {
                    "step": "File → Import → 选择 STEP 文件",
                    "iges": "File → Import → 选择 IGES 文件",
                    "brep": "File → Open → 选择 BRep 文件",
                    "stl": "File → Import → Mesh 格式",
                },
            },
            "solidworks": {
                "name": "SolidWorks",
                "primary_format": "step",
                "step_schema": "AP214",
                "secondary_format": "iges",
                "notes": [
                    "STEP AP214 最佳兼容性",
                    "导入后自动识别为实体",
                    "IGES 可能产生曲面而非实体",
                ],
                "import_settings": {
                    "step": "File → Open → STEP 文件",
                    "iges": "File → Open → IGES 文件",
                },
            },
        }

        return recommendations.get(target_software, recommendations["ug"])

    @staticmethod
    def _format_file_size(size_bytes: int) -> str:
        """格式化文件大小"""
        if size_bytes < 1024:
            return f"{size_bytes} B"
        elif size_bytes < 1024 * 1024:
            return f"{size_bytes / 1024:.1f} KB"
        else:
            return f"{size_bytes / (1024 * 1024):.1f} MB"