blender-vrm/vrm_mesh_tools/ops.py

# SPDX-License-Identifier: MIT
# VRM Mesh Tools - Operators for mesh separation and UV edge merging

from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional

import bpy
from bpy.types import Context, Operator

from .utils import (
    AUTO_CONVERT_MATERIALS,
    PROCESSING_ORDER,
    MaterialInfo,
    find_body_mesh,
    find_face_mesh,
    generate_unique_name,
    get_material_face_count,
    get_material_index,
    get_target_mesh_name,
    is_separable_material,
    isolate_material_faces,
    load_edge_uv_data,
    normalize_uv_pair,
    select_edges_by_uv,
)

if TYPE_CHECKING:
    from bpy.types import Object


@dataclass
class ProcessResult:
    """材质处理结果"""
    material_name: str
    mesh_name: str
    edges_selected: int
    edges_dissolved: int
    success: bool
    error_message: Optional[str] = None


@dataclass
class MergeOperationState:
    """操作状态，用于错误恢复"""
    original_mode: str
    original_active_object: Optional["Object"]
    original_selected_objects: list["Object"]
    processed_materials: list[str]


class VRM_TOOLS_OT_separate_non_skin_faces(Operator):
    """将Body网格中非皮肤材质的面分离为独立对象"""

    bl_idname = "vrm_tools.separate_non_skin_faces"
    bl_label = "Separate Non-Skin Faces"
    bl_description = "将Body网格中可分离材质（CLOTH和其他）的面合并分离为单个独立对象"
    bl_options: AbstractSet[str] = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context: Context) -> bool:
        """检查是否可以执行分离操作"""
        armature = cls._get_armature(context)
        if armature is None:
            return False

        body_mesh = find_body_mesh(armature)
        if body_mesh is None:
            return False

        # 检查是否有可分离材质
        mesh_data = body_mesh.data
        if mesh_data is None:
            return False

        for material in mesh_data.materials:
            if material is not None and is_separable_material(material.name):
                return True

        return False

    @classmethod
    def _get_armature(cls, context: Context) -> Optional["Object"]:
        """获取当前选中的Armature对象"""
        active_object = context.active_object
        if active_object is None:
            return None

        if active_object.type == "ARMATURE":
            return active_object

        # 如果选中的是网格，检查其父对象是否为Armature
        if active_object.parent is not None and active_object.parent.type == "ARMATURE":
            return active_object.parent

        return None

    def execute(self, context: Context) -> set[str]:
        """执行分离操作"""
        try:
            return self._execute_separation(context)
        except Exception as e:
            self.report({"ERROR"}, f"分离操作失败: {str(e)}")
            return {"CANCELLED"}

    def _execute_separation(self, context: Context) -> set[str]:
        """执行分离操作的核心逻辑"""
        armature = self._get_armature(context)
        if armature is None:
            self.report({"ERROR"}, "未找到Armature对象")
            return {"CANCELLED"}

        body_mesh = find_body_mesh(armature)
        if body_mesh is None:
            self.report({"ERROR"}, "未找到Body网格对象")
            return {"CANCELLED"}

        mesh_data = body_mesh.data
        if mesh_data is None:
            self.report({"ERROR"}, "Body网格数据无效")
            return {"CANCELLED"}

        # 收集可分离材质信息
        separable_materials: list[MaterialInfo] = []
        for i, material in enumerate(mesh_data.materials):
            if material is not None and is_separable_material(material.name):
                face_count = get_material_face_count(mesh_data, i)
                if face_count > 0:
                    separable_materials.append(
                        MaterialInfo(
                            name=material.name,
                            index=i,
                            face_count=face_count,
                            is_skin=False,
                        )
                    )

        if not separable_materials:
            self.report({"INFO"}, "没有找到可分离材质的面")
            return {"FINISHED"}

        # 获取现有对象名称集合
        existing_names = set(obj.name for obj in context.blend_data.objects)

        # 将所有可分离材质的面合并分离为单个对象
        material_indices = [mat_info.index for mat_info in separable_materials]
        new_obj = self._separate_all_material_faces(
            context, body_mesh, armature, material_indices, existing_names
        )

        if new_obj is not None:
            total_faces = sum(mat_info.face_count for mat_info in separable_materials)
            self.report(
                {"INFO"},
                f"成功分离了 {len(separable_materials)} 个材质的 {total_faces} 个面到 {new_obj.name}"
            )
            return {"FINISHED"}
        else:
            self.report({"WARNING"}, "没有成功分离任何面")
            return {"CANCELLED"}

    def _separate_all_material_faces(
        self,
        context: Context,
        body_obj: "Object",
        armature: "Object",
        material_indices: list[int],
        existing_names: set[str],
    ) -> Optional["Object"]:
        """将所有指定材质的面合并分离为单个新对象"""
        mesh_data = body_obj.data
        if mesh_data is None:
            return None

        # 确保在对象模式
        if context.mode != "OBJECT":
            bpy.ops.object.mode_set(mode="OBJECT")

        # 取消所有选择
        bpy.ops.object.select_all(action="DESELECT")

        # 选择Body网格并设为活动对象
        body_obj.select_set(True)
        context.view_layer.objects.active = body_obj

        # 进入编辑模式
        bpy.ops.object.mode_set(mode="EDIT")

        # 取消所有选择
        bpy.ops.mesh.select_all(action="DESELECT")

        # 切换到面选择模式
        bpy.ops.mesh.select_mode(type="FACE")

        # 返回对象模式以访问网格数据
        bpy.ops.object.mode_set(mode="OBJECT")

        # 选择所有使用可分离材质的面
        material_index_set = set(material_indices)
        for polygon in mesh_data.polygons:
            if polygon.material_index in material_index_set:
                polygon.select = True

        # 返回编辑模式执行分离
        bpy.ops.object.mode_set(mode="EDIT")

        # 执行分离操作
        bpy.ops.mesh.separate(type="SELECTED")

        # 返回对象模式
        bpy.ops.object.mode_set(mode="OBJECT")

        # 查找新创建的对象
        new_obj = self._find_new_separated_object(context, body_obj)
        if new_obj is None:
            return None

        # 使用固定名称"Body_Separated"，处理命名冲突
        base_name = "Body_Separated"
        unique_name = generate_unique_name(base_name, existing_names)
        new_obj.name = unique_name

        # 设置父子关系（保持与原Armature的关系）
        new_obj.parent = armature
        new_obj.parent_type = "OBJECT"

        return new_obj

    def _find_new_separated_object(
        self, context: Context, original_obj: "Object"
    ) -> Optional["Object"]:
        """查找分离操作后新创建的对象"""
        for obj in context.selected_objects:
            if obj != original_obj and obj.type == "MESH":
                return obj
        return None


class VRM_TOOLS_OT_merge_uv_edges(Operator):
    """根据UV坐标融并边，将三角面转四边面"""

    bl_idname = "vrm_tools.merge_uv_edges"
    bl_label = "Merge UV Edges"
    bl_description = "根据UV坐标选中并融并边，将三角面转换为四边面"
    bl_options: AbstractSet[str] = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context: Context) -> bool:
        """检查是否可以执行操作"""
        armature = cls._get_armature(context)
        if armature is None:
            return False

        body_mesh = find_body_mesh(armature)
        face_mesh = find_face_mesh(armature)

        return body_mesh is not None and face_mesh is not None

    @classmethod
    def _get_armature(cls, context: Context) -> Optional["Object"]:
        """获取当前选中的Armature对象"""
        active_object = context.active_object
        if active_object is None:
            return None

        if active_object.type == "ARMATURE":
            return active_object

        if active_object.parent is not None and active_object.parent.type == "ARMATURE":
            return active_object.parent

        return None

    def execute(self, context: Context) -> set[str]:
        """执行边融并操作"""
        state = self._save_state(context)

        try:
            return self._execute_merge(context, state)
        except Exception as e:
            self._restore_state(context, state)
            self.report({"ERROR"}, f"边融并操作失败: {str(e)}")
            return {"CANCELLED"}

    def _save_state(self, context: Context) -> MergeOperationState:
        """保存当前Blender状态"""
        original_mode = context.object.mode if context.object else "OBJECT"
        original_active = context.view_layer.objects.active
        original_selected = list(context.selected_objects)

        return MergeOperationState(
            original_mode=original_mode,
            original_active_object=original_active,
            original_selected_objects=original_selected,
            processed_materials=[],
        )

    def _restore_state(self, context: Context, state: MergeOperationState) -> None:
        """恢复原始Blender状态"""
        try:
            if context.object and context.object.mode != "OBJECT":
                bpy.ops.object.mode_set(mode="OBJECT")

            bpy.ops.object.select_all(action="DESELECT")

            for obj in state.original_selected_objects:
                if obj and obj.name in context.blend_data.objects:
                    obj.select_set(True)

            if state.original_active_object and \
               state.original_active_object.name in context.blend_data.objects:
                context.view_layer.objects.active = state.original_active_object

            if context.object and state.original_mode != "OBJECT":
                try:
                    bpy.ops.object.mode_set(mode=state.original_mode)
                except RuntimeError:
                    pass
        except Exception:
            pass

    def _execute_merge(
        self, context: Context, state: MergeOperationState
    ) -> set[str]:
        """执行边融并的核心逻辑"""
        armature = self._get_armature(context)
        if armature is None:
            self.report({"ERROR"}, "未找到Armature对象")
            return {"CANCELLED"}

        body_mesh = find_body_mesh(armature)
        face_mesh = find_face_mesh(armature)

        if body_mesh is None or face_mesh is None:
            self.report({"ERROR"}, "未找到Body或Face网格对象")
            return {"CANCELLED"}

        if context.object and context.object.mode != "OBJECT":
            bpy.ops.object.mode_set(mode="OBJECT")

        results: list[ProcessResult] = []
        total_edges_dissolved = 0
        materials_processed = 0
        materials_skipped = 0

        for material_name in PROCESSING_ORDER:
            target_mesh_name = get_target_mesh_name(material_name)
            if target_mesh_name is None:
                continue

            target_mesh = body_mesh if target_mesh_name == "Body" else face_mesh

            result = self._process_material(context, target_mesh, material_name)
            results.append(result)
            state.processed_materials.append(material_name)

            if result.success:
                materials_processed += 1
                total_edges_dissolved += result.edges_dissolved
            else:
                materials_skipped += 1
                if result.error_message:
                    self.report({"WARNING"}, result.error_message)

        if context.object and context.object.mode != "OBJECT":
            bpy.ops.object.mode_set(mode="OBJECT")

        if materials_processed > 0:
            self.report(
                {"INFO"},
                f"处理完成: {materials_processed} 个材质, "
                f"融并 {total_edges_dissolved} 条边"
            )
            return {"FINISHED"}
        else:
            self.report({"WARNING"}, f"没有成功处理任何材质 (跳过 {materials_skipped} 个)")
            return {"CANCELLED"}

    def _process_material(
        self,
        context: Context,
        mesh_obj: "Object",
        material_name: str,
    ) -> ProcessResult:
        """处理单个材质的边融并"""
        import bmesh

        mesh_name = mesh_obj.name if mesh_obj else "Unknown"

        if material_name in AUTO_CONVERT_MATERIALS:
            return self._process_material_auto_convert(context, mesh_obj, material_name)

        try:
            target_uvs = load_edge_uv_data(material_name)
            if target_uvs is None:
                return ProcessResult(
                    material_name=material_name,
                    mesh_name=mesh_name,
                    edges_selected=0,
                    edges_dissolved=0,
                    success=False,
                    error_message=f"跳过: 未找到材质 {material_name} 的JSON文件",
                )

            if len(target_uvs) == 0:
                return ProcessResult(
                    material_name=material_name,
                    mesh_name=mesh_name,
                    edges_selected=0,
                    edges_dissolved=0,
                    success=False,
                    error_message=f"跳过: 材质 {material_name} 的JSON文件中没有边数据",
                )

            if not isolate_material_faces(context, mesh_obj, material_name):
                return ProcessResult(
                    material_name=material_name,
                    mesh_name=mesh_name,
                    edges_selected=0,
                    edges_dissolved=0,
                    success=False,
                    error_message=f"跳过: 无法隔离材质 {material_name} 的面",
                )

            bpy.ops.mesh.select_all(action="DESELECT")

            mesh_data = mesh_obj.data
            bm = bmesh.from_edit_mesh(mesh_data)

            uv_layer = bm.loops.layers.uv.active
            if uv_layer is None:
                bpy.ops.mesh.reveal()
                bpy.ops.object.mode_set(mode="OBJECT")
                return ProcessResult(
                    material_name=material_name,
                    mesh_name=mesh_name,
                    edges_selected=0,
                    edges_dissolved=0,
                    success=False,
                    error_message=f"跳过: 网格 {mesh_name} 没有UV层",
                )

            edges_selected = select_edges_by_uv(bm, uv_layer, target_uvs)

            bmesh.update_edit_mesh(mesh_data)

            if edges_selected == 0:
                bpy.ops.mesh.reveal()
                bpy.ops.object.mode_set(mode="OBJECT")
                return ProcessResult(
                    material_name=material_name,
                    mesh_name=mesh_name,
                    edges_selected=0,
                    edges_dissolved=0,
                    success=False,
                    error_message=(
                        f"跳过: 材质 {material_name} 没有找到匹配的边 "
                        f"(目标边数: {len(target_uvs)})"
                    ),
                )

            edges_dissolved = 0
            try:
                bpy.ops.mesh.dissolve_edges()
                edges_dissolved = edges_selected
            except RuntimeError as e:
                self.report({"WARNING"}, f"材质 {material_name} 融并边时出现问题: {str(e)}")

            bpy.ops.mesh.reveal()

            bpy.ops.object.mode_set(mode="OBJECT")

            return ProcessResult(
                material_name=material_name,
                mesh_name=mesh_name,
                edges_selected=edges_selected,
                edges_dissolved=edges_dissolved,
                success=True,
                error_message=None,
            )

        except Exception as e:
            try:
                if context.object and context.object.mode == "EDIT":
                    bpy.ops.mesh.reveal()
                    bpy.ops.object.mode_set(mode="OBJECT")
            except Exception:
                pass

            return ProcessResult(
                material_name=material_name,
                mesh_name=mesh_name,
                edges_selected=0,
                edges_dissolved=0,
                success=False,
                error_message=f"处理材质 {material_name} 时发生错误: {str(e)}",
            )

    def _process_material_auto_convert(
        self,
        context: Context,
        mesh_obj: "Object",
        material_name: str,
    ) -> ProcessResult:
        """使用Blender自带的tris_convert_to_quads处理规律网格的材质"""
        mesh_name = mesh_obj.name if mesh_obj else "Unknown"

        try:
            if not isolate_material_faces(context, mesh_obj, material_name):
                return ProcessResult(
                    material_name=material_name,
                    mesh_name=mesh_name,
                    edges_selected=0,
                    edges_dissolved=0,
                    success=False,
                    error_message=f"跳过: 无法隔离材质 {material_name} 的面",
                )

            bpy.ops.mesh.select_mode(type="FACE")

            bpy.ops.mesh.select_all(action="SELECT")

            bpy.ops.mesh.tris_convert_to_quads()

            bpy.ops.mesh.reveal()

            bpy.ops.object.mode_set(mode="OBJECT")

            return ProcessResult(
                material_name=material_name,
                mesh_name=mesh_name,
                edges_selected=0,
                edges_dissolved=0,
                success=True,
                error_message=None,
            )

        except Exception as e:
            try:
                if context.object and context.object.mode == "EDIT":
                    bpy.ops.mesh.reveal()
                    bpy.ops.object.mode_set(mode="OBJECT")
            except Exception:
                pass

            return ProcessResult(
                material_name=material_name,
                mesh_name=mesh_name,
                edges_selected=0,
                edges_dissolved=0,
                success=False,
                error_message=f"处理材质 {material_name} 时发生错误: {str(e)}",
            )