# src/core/mold_generator.py
from pathlib import Path
from typing import Dict, List, Any, Tuple, Optional
import numpy as np
from OCC.Core.BRepOffsetAPI import BRepOffsetAPI_MakeThickSolid
from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Cut, BRepAlgoAPI_Fuse
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_MakeFace, BRepBuilderAPI_Transform
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeBox
from OCC.Core.Geom import Geom_Plane
from OCC.Core.gp import gp_Pln, gp_Dir, gp_Pnt, gp_Vec, gp_Trsf
from OCC.Core.TopTools import TopTools_ListOfShape
from OCC.Core.TopoDS import TopoDS_Face, TopoDS_Shape
from OCC.Core.BRep import BRep_Tool
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.BRepExtrema import BRepExtrema_DistShapeShape
from OCC.Core.GProp import GProp_GProps
from OCC.Core.BRepGProp import brepgprop

from models.schemas import create_mold_cavity_data, create_mold_key_info
from utils.logger import get_logger

logger = get_logger(__name__)


class MoldCavityGenerator:
    """模具型腔生成器 - 基于产品模型生成Cavity和Core"""

    def __init__(self, shrinkage_rate: float = 0.005, draft_angle: float = 2.0,
                 material_density: float = 1.05):
        """
        初始化模具生成器

        Args:
            shrinkage_rate: 收缩率（默认0.5% for ABS）
            draft_angle: 拔模角（默认2度）
            material_density: 材料密度 g/cm³（默认1.05 for ABS）
        """
        self.shrinkage_rate = shrinkage_rate
        self.draft_angle = draft_angle  # 度
        self.material_density = material_density  # g/cm³

        # 常用塑料材料密度（g/cm³）
        self.material_densities = {
            "ABS": 1.05,
            "PP": 0.90,
            "PC": 1.20,
            "PE": 0.95,
            "PS": 1.05,
            "PA": 1.14,
            "POM": 1.42,
            "PMMA": 1.18
        }

        # 分型面检测参数
        self.parting_line_tolerance = 0.1
        self.max_draft_angle = 5.0

    def set_material(self, material: str):
        """设置产品材料"""
        if material in self.material_densities:
            self.material_density = self.material_densities[material]
            logger.info(f"材料设置为 {material}, 密度: {self.material_density} g/cm³")
        else:
            logger.warning(f"未知材料 {material}, 使用默认密度 {self.material_density} g/cm³")

    def generate_mold_cavities(self, product_shape: Any) -> Dict[str, Any]:
        """
        从产品的3D模型生成型腔和型芯

        Returns:
            {
                "cavity": cavity_shape,      # 型腔（产品外部）
                "core": core_shape,          # 型芯（产品内部）
                "parting_surface": parting_surface,  # 分型面
                "parting_line": parting_line         # 分型线
            }
        """
        logger.info("开始生成模具型腔...")

        try:
            # Step 1: 分析产品几何
            analysis = self._analyze_product_geometry(product_shape)

            # Step 2: 检测分型面和分型线
            parting_surface, parting_line = self._detect_parting_surface(
                product_shape, analysis
            )

            # Step 3: 应用收缩率补偿
            scaled_shape = self._apply_shrinkage_compensation(product_shape)

            # Step 4: 添加拔模角
            drafted_shape = self._apply_draft_angles(scaled_shape, parting_surface)

            # Step 5: 分离型腔和型芯
            cavity, core = self._split_cavity_core(drafted_shape, parting_surface)

            logger.info("模具型腔生成完成")

            return {
                "cavity": cavity,
                "core": core,
                "parting_surface": parting_surface,
                "parting_line": parting_line,
                "analysis": analysis
            }

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

    def generate_detailed_cavity_json(self, cavity_data: Dict) -> Dict[str, Any]:
        """
        生成详细的型腔三维JSON数据

        Returns:
            包含完整几何信息的JSON结构
        """
        cavity = cavity_data["cavity"]
        core = cavity_data["core"]
        parting_surface = cavity_data["parting_surface"]
        analysis = cavity_data["analysis"]

        # 提取型腔几何数据
        cavity_geometry = self._extract_shape_geometry(cavity, "cavity")
        core_geometry = self._extract_shape_geometry(core, "core")

        # 提取分型面数据
        parting_geometry = self._extract_parting_surface_geometry(
            parting_surface
        )

        detailed_json = {
            "metadata": {
                "version": "2.0",
                "generated_at": str(np.datetime64('now')),
                "shrinkage_rate": self.shrinkage_rate,
                "draft_angle": self.draft_angle,
                "unit": "mm"
            },
            "product_analysis": {
                "bounding_box": analysis.get("bounding_box", {}),  # 使用get方法
                "volume": analysis.get("volume", 0),               # 使用get方法
                "surface_area": analysis.get("surface_area", 0),   # 使用get方法
                "center_of_mass": analysis.get("center_of_mass", [0, 0, 0])  # 使用get方法
            },
            "mold_cavities": {
                "cavity": cavity_geometry,
                "core": core_geometry
            },
            "parting_surface": parting_geometry,
            "manufacturing_info": {
                "estimated_mold_size": self._calculate_mold_size(analysis),
                "estimated_clamping_force": self._calculate_clamping_force(analysis),
                "recommended_material": self._get_recommended_material()
            }
        }

        return detailed_json

    def generate_cavity_key_info(self, cavity_data: Dict) -> Dict[str, Any]:
        """
        生成模具型腔的关键信息

        Returns:
            关键参数摘要
        """
        analysis = cavity_data["analysis"]

        key_info = {
            "mold_parameters": {
                "shrinkage_rate": f"{self.shrinkage_rate * 100:.2f}%",
                "draft_angle": f"{self.draft_angle}°",
                "parting_line_length": self._calculate_parting_line_length(
                    cavity_data["parting_line"]
                ),
                "cavity_depth": analysis.get("bounding_box", {}).get("dimensions", [0, 0, 0])[2]
            },
            "geometric_characteristics": {
                "product_volume": f"{analysis.get('volume', 0) / 1000:.2f} cm³",
                "product_weight": self._calculate_product_weight(analysis),
                "wall_thickness_range": self._estimate_wall_thickness(analysis),
                "complexity_score": self._calculate_complexity_score(analysis)
            },
            "manufacturing_requirements": {
                "cavity_material": "Aluminum Alloy 7075",
                "hardness": "HRC 30-35",
                "surface_finish": "SPI A2",
                "estimated_cycle_time": self._estimate_cycle_time(analysis),
                "recommended_injection_pressure": "80-120 MPa"
            },
            "quality_considerations": {
                "potential_weld_lines": self._identify_weld_line_risk(analysis),
                "sink_mark_areas": self._identify_sink_mark_risk(analysis),
                "warpage_risk": self._assess_warpage_risk(analysis)
            }
        }

        return key_info

    # ==================== 内部方法 ====================

    def _analyze_product_geometry(self, shape: Any) -> Dict[str, Any]:
        """分析产品几何属性"""
        # 计算体积属性
        volume_props = GProp_GProps()
        brepgprop.VolumeProperties(shape, volume_props)

        # 计算表面积属性
        surface_props = GProp_GProps()
        brepgprop.SurfaceProperties(shape, surface_props)

        # 计算边界框
        from OCC.Core.Bnd import Bnd_Box
        from OCC.Core.BRepBndLib import brepbndlib

        bbox = Bnd_Box()
        brepbndlib.Add(shape, bbox)
        xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()

        return {
            "volume": volume_props.Mass(),
            "surface_area": surface_props.Mass(),
            "center_of_mass": [
                volume_props.CentreOfMass().X(),
                volume_props.CentreOfMass().Y(),
                volume_props.CentreOfMass().Z()
            ],
            "bounding_box": {
                "min": [xmin, ymin, zmin],
                "max": [xmax, ymax, zmax],
                "dimensions": [xmax - xmin, ymax - ymin, zmax - zmin],
                "center": [(xmin + xmax) / 2, (ymin + ymax) / 2, (zmin + zmax) / 2]
            },
            "inertia_matrix": self._get_inertia_matrix(volume_props)
        }

    def _detect_parting_surface(self, shape: Any, analysis: Dict) -> Tuple[Any, List]:
        """检测分型面和分型线"""
        # 简化的分型面检测：基于Z方向的最高点和最低点
        bbox = analysis["bounding_box"]
        center_z = bbox["center"][2]

        # 创建分型面（XY平面）
        parting_plane = gp_Pln(
            gp_Pnt(0, 0, center_z),
            gp_Dir(0, 0, 1)
        )
        parting_surface = BRepBuilderAPI_MakeFace(
            parting_plane,
            bbox["min"][0] - 10, bbox["max"][0] + 10,
            bbox["min"][1] - 10, bbox["max"][1] + 10
        ).Face()

        # 分型线（简化）
        parting_line = [
            [bbox["min"][0], bbox["min"][1], center_z],
            [bbox["max"][0], bbox["min"][1], center_z],
            [bbox["max"][0], bbox["max"][1], center_z],
            [bbox["min"][0], bbox["max"][1], center_z],
            [bbox["min"][0], bbox["min"][1], center_z]
        ]

        return parting_surface, parting_line

    def _apply_shrinkage_compensation(self, shape: Any) -> Any:
        """应用收缩率补偿（放大模型）"""
        scale_factor = 1.0 + self.shrinkage_rate

        # 创建缩放变换
        trsf = gp_Trsf()
        trsf.SetScale(gp_Pnt(0, 0, 0), scale_factor)

        from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
        scaled_shape = BRepBuilderAPI_Transform(shape, trsf, True).Shape()

        return scaled_shape

    def _apply_draft_angles(self, shape: Any, parting_surface: Any) -> Any:
        """添加拔模角（简化实现）"""
        # 实际实现需要复杂的拔模面处理
        # 这里返回原始形状（假设已在CAD中处理）
        logger.warning("拔模角处理为简化实现，建议在设计阶段处理")
        return shape

    def _split_cavity_core(self, shape: Any, parting_surface: Any) -> Tuple[Any, Any]:
        """分离型腔和型芯
        
        型腔(Cavity): 模具中形成产品外表面的部分，是产品形状的负形
        型芯(Core): 模具中形成产品内表面的部分，是产品形状的正形
        """
        try:
            from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeBox
            from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Cut
            from OCC.Core.TopExp import TopExp_Explorer
            from OCC.Core.TopAbs import TopAbs_SOLID
            
            # 获取产品边界框
            from OCC.Core.Bnd import Bnd_Box
            from OCC.Core.BRepBndLib import brepbndlib
            bbox = Bnd_Box()
            brepbndlib.Add(shape, bbox)
            xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()
            
            # 计算模具块尺寸（比产品大一定余量）
            margin = 20  # mm
            mold_xmin = xmin - margin
            mold_ymin = ymin - margin
            mold_zmin = zmin - margin
            mold_xmax = xmax + margin
            mold_ymax = ymax + margin
            mold_zmax = zmax + margin
            
            # 创建模具块
            mold_block = BRepPrimAPI_MakeBox(
                gp_Pnt(mold_xmin, mold_ymin, mold_zmin),
                gp_Pnt(mold_xmax, mold_ymax, mold_zmax)
            ).Shape()
            
            # 型腔 = 模具块 - 产品（布尔减法）
            cavity_operation = BRepAlgoAPI_Cut(mold_block, shape)
            if cavity_operation.IsDone():
                cavity = cavity_operation.Shape()
                logger.info("型腔生成成功（模具块减去产品）")
            else:
                logger.warning("型腔布尔运算失败，使用原始形状")
                cavity = mold_block
            
            # 型芯 = 产品形状本身（收缩补偿后）
            core = shape
            logger.info("型芯 = 产品形状")
            
            return cavity, core

        except Exception as e:
            logger.error(f"型腔分离失败: {e}")
            return shape, shape

    def _extract_shape_geometry(self, shape: Any, shape_type: str) -> Dict[str, Any]:
        """提取形状几何数据为JSON格式"""
        try:
            # 网格化
            mesh = BRepMesh_IncrementalMesh(shape, 0.1)
            mesh.Perform()

            # 提取顶点和面
            from OCC.Core.TopExp import TopExp_Explorer
            from OCC.Core.TopAbs import TopAbs_FACE
            from OCC.Core.BRep import BRep_Tool
            from OCC.Core.Poly import Poly_Triangulation
            from OCC.Core.TopLoc import TopLoc_Location

            vertices = []
            faces = []

            explorer = TopExp_Explorer(shape, TopAbs_FACE)
            vertex_index = 0

            while explorer.More():
                # 使用 explorer.Current() 直接获取面
                face = explorer.Current()
                location = TopLoc_Location()
                triangulation = BRep_Tool.Triangulation(face, location)

                if triangulation:
                    # 提取顶点
                    nb_nodes = triangulation.NbNodes()
                    for i in range(1, nb_nodes + 1):
                        node = triangulation.Node(i)
                        # 应用位置变换
                        transformed = node.Transformed(location.Transformation())
                        vertices.extend([
                            float(transformed.X()),
                            float(transformed.Y()),
                            float(transformed.Z())
                        ])

                    # 提取三角形面
                    nb_triangles = triangulation.NbTriangles()
                    for i in range(1, nb_triangles + 1):
                        triangle = triangulation.Triangle(i)
                        # 三角形顶点索引需要加上之前的顶点数量
                        idx1 = triangle.Value(1) + vertex_index - 1
                        idx2 = triangle.Value(2) + vertex_index - 1
                        idx3 = triangle.Value(3) + vertex_index - 1
                        faces.extend([int(idx1), int(idx2), int(idx3)])

                    vertex_index += nb_nodes

                explorer.Next()

            vertex_count = len(vertices) // 3
            face_count = len(faces) // 3

            return {
                "type": shape_type,
                "vertices": vertices,
                "faces": faces,
                "vertex_count": vertex_count,
                "face_count": face_count,
                "triangulation": "BRepMesh三角化"
            }

        except Exception as e:
            logger.error(f"{shape_type}几何提取失败: {e}")
            return {
                "type": shape_type,
                "vertices": [],
                "faces": [],
                "vertex_count": 0,
                "face_count": 0,
                "triangulation": f"提取失败: {str(e)}"
            }

    def _extract_parting_surface_geometry(self, surface: Any) -> Dict[str, Any]:
        """提取分型面几何数据"""
        # 尝试从surface获取边界信息，失败则使用默认值
        try:
            from OCC.Core.BRepAdaptor import BRepAdaptor_Surface
            adaptor = BRepAdaptor_Surface(surface)
            u_min, u_max = adaptor.FirstUParameter(), adaptor.LastUParameter()
            v_min, v_max = adaptor.FirstVParameter(), adaptor.LastVParameter()

            bounds = {
                "u_range": [float(u_min), float(u_max)],
                "v_range": [float(v_min), float(v_max)]
            }
        except Exception as e:
            logger.warning(f"分型面边界提取失败，使用默认值: {e}")
            bounds = {
                "u_range": [-200, 200],
                "v_range": [-200, 200]
            }

        # 分型面是水平面，法向量为 [0, 0, 1]，原点在 Z 轴中心
        return {
            "type": "plane",
            "normal": [0, 0, 1],
            "origin": [0, 0, 0],
            "bounds": bounds
        }

        return {
            "type": "plane",
            "normal": [0, 0, 1],
            "origin": [0, 0, 0],
            "bounds": bounds
        }

    def _calculate_mold_size(self, analysis: Dict) -> Dict[str, float]:
        """估算模具尺寸"""
        product_bbox = analysis["bounding_box"]["dimensions"]

        # 模具通常比产品大20-50mm
        margin = 30  # mm

        return {
            "length": product_bbox[0] + 2 * margin,
            "width": product_bbox[1] + 2 * margin,
            "height": product_bbox[2] + 2 * margin + 100,  # 增加100mm用于模架
            "margin": margin
        }

    def _calculate_clamping_force(self, analysis: Dict) -> str:
        """估算锁模力"""
        volume_cm3 = analysis.get("volume", 0) / 1000  # mm³ → cm³

        # 经验公式: 锁模力 ≈ 投影面积 × 压力 × 安全系数
        # 简化估算
        if volume_cm3 < 10:
            return "50-100 吨"
        elif volume_cm3 < 100:
            return "150-300 吨"
        elif volume_cm3 < 500:
            return "400-600 吨"
        else:
            return "800+ 吨"

    def _get_recommended_material(self) -> str:
        """推荐模具材料"""
        # 根据产品产量推荐模具材料
        # 小批量 (<5000件): 铝合金
        # 中批量 (5000-50000件): P20钢
        # 大批量 (>50000件): H13钢
        return "Aluminum Alloy 7075 (铝合金模具)"

    def _calculate_product_weight(self, analysis: Dict) -> str:
        """计算产品重量（使用当前材料密度）"""
        volume_cm3 = analysis.get("volume", 0) / 1000
        weight_g = volume_cm3 * self.material_density
        return f"{weight_g:.2f} g"

    def _estimate_wall_thickness(self, analysis: Dict) -> str:
        """估算壁厚范围"""
        volume = analysis.get("volume", 0)
        surface_area = analysis.get("surface_area", 0)

        if surface_area > 0 and volume > 0:
            avg_thickness = (volume / surface_area) * 0.6
            return f"{avg_thickness * 0.7:.2f} - {avg_thickness * 1.3:.2f} mm"
        elif volume > 0:
            # 如果没有surface_area，基于体积估算
            bbox_dims = analysis.get("bounding_box", {}).get("dimensions", [1, 1, 1])
            bbox_volume = bbox_dims[0] * bbox_dims[1] * bbox_dims[2]
            if bbox_volume > 0:
                efficiency = volume / bbox_volume
                avg_thickness = (bbox_dims[0] + bbox_dims[1]) / 2 * efficiency
                return f"{avg_thickness * 0.7:.2f} - {avg_thickness * 1.3:.2f} mm"

        return "2.0 - 4.0 mm (默认)"

    def _calculate_complexity_score(self, analysis: Dict) -> float:
        """计算复杂度评分（0-10）"""
        # 基于体积、表面积比、边界框等
        volume = analysis.get("volume", 0)
        surface_area = analysis.get("surface_area", 0)

        if surface_area > 0 and volume > 0:
            thickness_ratio = (volume / surface_area) * 0.6
            complexity = min(thickness_ratio / 5.0, 10.0)
            return round(complexity, 1)
        elif volume > 0:
            # 如果没有surface_area，基于拓扑复杂度评分
            bbox_dims = analysis.get("bounding_box", {}).get("dimensions", [100, 100, 100])
            bbox_volume = bbox_dims[0] * bbox_dims[1] * bbox_dims[2]
            if bbox_volume > 0:
                volume_ratio = volume / bbox_volume
                complexity = (1.0 - volume_ratio) * 10
                return round(min(max(complexity, 0), 10), 1)

        return 5.0

    def _estimate_cycle_time(self, analysis: Dict) -> str:
        """估算成型周期"""
        volume_cm3 = analysis.get("volume", 0) / 1000

        if volume_cm3 < 10:
            return "15-25 秒"
        elif volume_cm3 < 50:
            return "25-40 秒"
        elif volume_cm3 < 200:
            return "40-60 秒"
        else:
            return "60-90 秒"

    def _identify_weld_line_risk(self, analysis: Dict) -> str:
        """识别熔接痕风险"""
        # 基于几何复杂度判断
        complexity = self._calculate_complexity_score(analysis)

        if complexity > 7:
            return "高 - 建议优化浇口位置"
        elif complexity > 4:
            return "中 - 需仿真验证"
        else:
            return "低"

    def _identify_sink_mark_risk(self, analysis: Dict) -> str:
        """识别缩痕风险"""
        thickness = self._estimate_wall_thickness(analysis)
        # 简化的风险评估
        return "中 - 建议壁厚均匀性检查"

    def _assess_warpage_risk(self, analysis: Dict) -> str:
        """评估翘曲风险"""
        bbox = analysis.get("bounding_box", {}).get("dimensions", [1, 1, 1])
        aspect_ratio = max(bbox) / min(bbox)

        if aspect_ratio > 5:
            return "高 - 建议增加加强筋"
        elif aspect_ratio > 3:
            return "中 - 需优化冷却"
        else:
            return "低"

    def _get_inertia_matrix(self, props: GProp_GProps) -> List[List[float]]:
        """获取惯性矩阵"""
        inertia = props.MatrixOfInertia()
        return [
            [inertia.Value(1, 1), inertia.Value(1, 2), inertia.Value(1, 3)],
            [inertia.Value(2, 1), inertia.Value(2, 2), inertia.Value(2, 3)],
            [inertia.Value(3, 1), inertia.Value(3, 2), inertia.Value(3, 3)]
        ]

    def _calculate_parting_line_length(self, parting_line: List) -> float:
        """计算分型线长度"""
        # 简化的长度计算
        return 250.0  # mm