import math
import os
import sys

import bpy
import numpy as np
from bpy.types import Action, Armature, Mesh, Object

# isort: split
import bmesh
import mathutils


class HiddenPrints:
    def __init__(self, enable=True, suppress_err=False):
        self.enable = enable
        self.suppress_err = suppress_err

    def __enter__(self):
        if not self.enable:
            return
        self._original_stdout = sys.stdout
        sys.stdout = open(os.devnull, "w")
        if self.suppress_err:
            self._original_stderr = sys.stderr
            sys.stderr = open(os.devnull, "w")

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self.enable:
            return
        sys.stdout.close()
        sys.stdout = self._original_stdout
        if self.suppress_err:
            sys.stderr.close()
            sys.stderr = self._original_stderr


USE_WORLD_COORDINATES = False


class Mode:
    def __init__(self, mode_name="EDIT", active_obj: Object = None):
        self.mode = mode_name
        self.active = active_obj
        self.pre_active = None
        self.pre_mode = "OBJECT"

    def __enter__(self):
        self.pre_active = bpy.context.view_layer.objects.active
        if self.pre_active is not None:
            self.pre_mode = bpy.context.object.mode
        bpy.context.view_layer.objects.active = self.active
        bpy.ops.object.mode_set(mode=self.mode)
        return self.active

    def __exit__(self, exc_type, exc_val, exc_tb):
        bpy.ops.object.mode_set(mode=self.pre_mode)
        bpy.context.view_layer.objects.active = self.pre_active


def reset():
    bpy.ops.wm.read_factory_settings(use_empty=True)


def remove_all(delete_actions=True):
    for obj in bpy.data.objects.values():
        bpy.data.objects.remove(obj, do_unlink=True)
    bpy.ops.outliner.orphans_purge(do_recursive=True)
    if delete_actions:
        for action in bpy.data.actions:
            bpy.data.actions.remove(action, do_unlink=True)


def remove_empty():
    childless_empties = [e for e in bpy.data.objects if e.type.startswith("EMPTY") and not e.children]
    bpy.data.batch_remove(childless_empties)


def load_file(filepath: str, *args, **kwargs) -> "list[Object]":
    old_objs = set(bpy.context.scene.objects)
    if filepath.endswith(".glb"):
        bpy.ops.import_scene.gltf(filepath=filepath, *args, **kwargs)
    elif filepath.endswith(".fbx"):
        bpy.ops.import_scene.fbx(filepath=filepath, *args, **kwargs)
    elif filepath.endswith(".obj"):
        bpy.ops.wm.obj_import(filepath=filepath, *args, **kwargs)
    elif filepath.endswith(".ply"):
        bpy.ops.wm.ply_import(filepath=filepath, *args, **kwargs)
    else:
        raise RuntimeError(f"Invalid input file: {filepath}")
    imported_objs = set(bpy.context.scene.objects) - old_objs
    imported_objs = sorted(imported_objs, key=lambda x: x.name)
    print("Imported:", imported_objs)
    return imported_objs


def select_all():
    bpy.ops.object.select_all(action="SELECT")


def deselect():
    bpy.ops.object.select_all(action="DESELECT")


def select_objs(obj_list: "list[Object]" = None, deselect_first=False):
    if not obj_list:
        obj_list = bpy.context.scene.objects
    if deselect_first:
        deselect()
    for obj in obj_list:
        obj.select_set(True)


def select_mesh(obj_list: "list[Object]" = None, all=True, deselect_first=False):
    if not obj_list:
        obj_list = bpy.context.scene.objects
    if deselect_first:
        deselect()
    for obj in obj_list:
        if obj.type == "MESH":
            if all:
                obj.select_set(True)
            else:
                break


class Select:
    """
    Deselecting before and after selecting the specified objects.
    """

    def __init__(self, objs: "Object | list[Object]" = None):
        self.objs = (objs,) if isinstance(objs, Object) else objs
        self.objs: "tuple[Object]" = tuple(self.objs)

    def __enter__(self):
        select_objs(self.objs, deselect_first=True)
        return self.objs

    def __exit__(self, exc_type, exc_val, exc_tb):
        deselect()


def get_type_objs(obj_list: "list[Object]" = None, type="MESH", sort=True) -> "list[Object]":
    if not obj_list:
        obj_list = bpy.context.scene.objects
    type_obj_list = [obj for obj in obj_list if obj.type == type]
    if sort:
        type_obj_list = sorted(type_obj_list, key=lambda x: x.name)
    return type_obj_list


def get_all_mesh_obj(obj_list: "list[Object]" = None):
    return get_type_objs(obj_list, "MESH")


def get_all_armature_obj(obj_list: "list[Object]" = None):
    return get_type_objs(obj_list, "ARMATURE")


def get_armature_obj(obj_list: "list[Object]" = None) -> Object:
    if not obj_list:
        obj_list = bpy.context.scene.objects
    for obj in obj_list:
        if obj.type == "ARMATURE":
            return obj


def get_rest_bones(armature_obj: Object):
    if armature_obj is None:
        return None, None, None
    rest_bones = []
    rest_bones_tail = []
    bones_idx_dict: "dict[str, int]" = {}
    armature_data: Armature = armature_obj.data
    for i, bone in enumerate(armature_data.bones):
        pos = bone.head_local
        pos_tail = bone.tail_local
        if USE_WORLD_COORDINATES:
            pos = armature_obj.matrix_world @ pos
            pos_tail = armature_obj.matrix_world @ pos_tail
        rest_bones.append(pos)
        rest_bones_tail.append(pos_tail)
        bones_idx_dict[bone.name] = i
    rest_bones = np.stack(rest_bones, axis=0)
    rest_bones_tail = np.stack(rest_bones_tail, axis=0)
    return rest_bones, rest_bones_tail, bones_idx_dict


def mesh_quads2tris(obj_list: "list[Object]" = None):
    if not obj_list:
        obj_list = bpy.context.scene.objects
    for obj in obj_list:
        if obj.type == "MESH":
            with Mode("EDIT", obj):
                bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")