# ##### BEGIN GPL LICENSE BLOCK #####
#
#  Copyright (C) 2022 VFX Grace - All Rights Reserved
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
# ##### END GPL LICENSE BLOCK #####

import os
import bpy
from bl_ui.space_toolsystem_common import ToolSelectPanelHelper, ToolDef, _icon_cache
from bl_ui.space_toolsystem_toolbar import VIEW3D_PT_tools_active, _defs_particle
from bpy.types import WorkSpaceTool, Panel, Operator
from gpu_extras.batch import batch_for_shader
from mathutils import Matrix
import numpy as np

from .particle_hook import particle_hook
from .draw_hair import show_hair_color, update_hair_color, ColorUpdateFlag
from .utils import *
from .props import *
from . import pie_menu

CIRCLE_WIRE_2D_POINTS = get_circle_wire_2d(0, 0, 40)
CIRCLE_WIRE_3D_POINTS = [(v[0], v[1], 0) for v in CIRCLE_WIRE_2D_POINTS]
BRUSH3D_IDNAME_PREFIX = "hair_brush_3d"

vertex_shader = '''
    uniform mat4 ModelViewProjectionMatrix;

    in vec3 pos;

    void main()
    {
        gl_Position = ModelViewProjectionMatrix * vec4(pos, 1.0);
    }
'''

fragment_shader = '''
    void main()
    {             
    }
'''

G_on_mesh = False
G_loc_world = Vector()
G_nor_world = Vector()
G_rad_3d = 0.0
G_brush_size = 0.0
G_pack_source_circle = False
G_source_circle = []
G_show_clone_circle = False
addon_name = get_addon_name()


def gpu_draw_shader(shader_name, type, color, coords, indices=None, line_width=1):
    shader = gpu.shader.from_builtin(shader_name)
    gl_blend_alpha(True)
    gl_line_smooth(True)
    gl_line_width(line_width)
    batch = batch_for_shader(shader, type, {"pos": coords}, indices=indices)
    shader.bind()
    shader.uniform_float("color", color)
    batch.draw(shader)
    gl_line_width(1)
    gl_line_smooth(False)
    gl_blend_alpha(False)


def draw_circle_cursor_2d(xy, rad):
    with gpu.matrix.push_pop():
        gpu.matrix.translate(xy)
        gpu.matrix.scale_uniform(rad)
        gpu_draw_shader('2D_UNIFORM_COLOR', 'LINE_STRIP', (1, 1, 1, 0.6), CIRCLE_WIRE_2D_POINTS, None, 2)


def draw_circle_cursor_3d(context, obmat, loc, nor, rad, clone_source=None):
    global G_pack_source_circle, G_source_circle, G_show_clone_circle
    with gpu.matrix.push_pop():
        window = context.window
        region = context.region
        rv3d = context.region_data
        obmat_ = obmat @ Matrix.Translation(loc)
        z_axis = Vector((0, 0, 1))
        rotmat = z_axis.rotation_difference(nor).to_matrix().to_4x4()
        gl_viewport(region.x, region.y, region.width, region.height)
        gpu.matrix.load_identity()
        gpu.matrix.push_projection()
        gpu.matrix.load_projection_matrix(rv3d.window_matrix)
        gpu.matrix.load_matrix(rv3d.view_matrix)

        if G_pack_source_circle:
            G_source_circle.clear()
            G_pack_source_circle = False
            for point in CIRCLE_WIRE_3D_POINTS:
                G_source_circle.append(obmat_ @ rotmat @ (Vector(point) * rad))
        if G_show_clone_circle and G_source_circle:
            gpu_draw_shader('3D_UNIFORM_COLOR', 'LINE_STRIP', (1, 0, 0, 0.7), G_source_circle, None, 2)

        if clone_source:
            gpu_draw_shader('3D_UNIFORM_COLOR', 'LINES', (0, 1, 0, 0.7), clone_source[0], clone_source[1])

        gpu.matrix.push()
        gpu.matrix.multiply_matrix(obmat_)
        gpu.matrix.multiply_matrix(rotmat)
        gpu.matrix.scale_uniform(rad)
        gpu_draw_shader('3D_UNIFORM_COLOR', 'LINE_STRIP', (1, 1, 0, 0.7), CIRCLE_WIRE_3D_POINTS, None, 2)
        gpu_draw_shader('3D_UNIFORM_COLOR', 'LINES', (1, 1, 0, 0.7), [(0, 0, 0), (0, 0, 2)], None, 2)
        gpu.matrix.pop()

        gpu.matrix.pop_projection()
        gl_viewport(0, 0, window.width, window.height)


class EditArgsFlag:
    USE_3D = 1
    ON_MESH = 2
    FOR_RENDER = 4
    USE_SCALE = 8


@classmethod
def brush3d_draw_cls(cls, layout, context, detect_layout=True, scale_y=1.75):
    global G_show_clone_circle
    tool_active_id = getattr(ToolSelectPanelHelper._tool_active_from_context(context, context.space_data.type),
                             "idname", None)
    if tool_active_id == "hair_brush_3d.Clone":
        G_show_clone_circle = True
    else:
        G_show_clone_circle = False

    cls.draw_cls__(layout, context, detect_layout=detect_layout, scale_y=scale_y)


@classmethod
def brush3d_tools_from_context(cls, context, mode=None):
    if mode is None:
        mode = context.mode
    for tools in (cls._tools[None], cls._tools.get(mode, ())):
        for item in tools:
            if not (type(item) is ToolDef) and callable(item):
                if (mode == 'PARTICLE' and context.preferences.addons[addon_name].preferences.collapse_intern
                        and item == _defs_particle.generate_from_brushes):
                    yield item(context)
                else:
                    yield from item(context)
            else:
                yield item


def patch_VIEW3D_PT_tools_active():
    VIEW3D_PT_tools_active.draw_cls__ = VIEW3D_PT_tools_active.draw_cls
    VIEW3D_PT_tools_active.tools_from_context__ = VIEW3D_PT_tools_active.tools_from_context
    VIEW3D_PT_tools_active.draw_cls = brush3d_draw_cls
    VIEW3D_PT_tools_active.tools_from_context = brush3d_tools_from_context


def restore_VIEW3D_PT_tools_active():
    VIEW3D_PT_tools_active.draw_cls = VIEW3D_PT_tools_active.draw_cls__
    VIEW3D_PT_tools_active.tools_from_context = VIEW3D_PT_tools_active.tools_from_context__
    del VIEW3D_PT_tools_active.draw_cls__
    del VIEW3D_PT_tools_active.tools_from_context__


class BrushType:
    __BRUSH_TYPES = {}

    @classmethod
    def add(cls, brush_cls):
        if brush_cls.bl_idname not in cls.__BRUSH_TYPES:
            cls.__BRUSH_TYPES[brush_cls.bl_idname] = brush_cls
            brush_cls.brush_type = len(cls.__BRUSH_TYPES) - 1

    @classmethod
    def get(cls, brush_idname: str) -> 'HairBrush3D':
        return cls.__BRUSH_TYPES[brush_idname]

    @classmethod
    def contains(cls, brush_idname: str):
        return brush_idname in cls.__BRUSH_TYPES


class HairBrush3D(WorkSpaceTool):
    bl_space_type = 'VIEW_3D'
    bl_context_mode = 'PARTICLE'
    bl_widget = None
    bl_keymap = (
        ("particle.brush3d_edit", {"type": 'LEFTMOUSE', "value": 'PRESS'}, {"properties": []}),
        ("particle.brush3d_edit", {"type": 'LEFTMOUSE', "value": 'PRESS', "shift": True}, {"properties": []}),)

    builtin_brush_id = ""
    support_2d = True
    support_3d = True

    @classmethod
    def can_cancel(cls, hb3):
        return hb3.use_3d_brush and not hb3.auto_switch_mode and not G_on_mesh

    @classmethod
    def can_to_builtin(cls, hb3):
        if not cls.builtin_brush_id:
            return False
        return not hb3.use_3d_brush or (hb3.auto_switch_mode and not G_on_mesh)

    @classmethod
    def can_update_3d_cursor(cls, hb3):
        return hb3.use_3d_brush

    @staticmethod
    def poll(obj):
        def is_psys_shown_hair(obj, psys):
            for mod in obj.modifiers:
                if mod.type == 'PARTICLE_SYSTEM' and mod.particle_system == psys:
                    return mod.show_viewport and psys.settings.type == 'HAIR'
            return False

        if (obj and obj.hide_viewport is False and obj.particle_systems.active and
                is_psys_shown_hair(obj, obj.particle_systems.active)):
            return True
        return False

    @staticmethod
    def draw_common_options(context, layout, draw_strenght=True, draw_adjust=True):
        pe = context.tool_settings.particle_edit
        hb3 = context.scene.hair_brush_3d
        layout.prop(pe.brush, "size", slider=True)
        if draw_strenght:
            layout.prop(pe.brush, "strength", slider=True)
        if draw_adjust:
            layout.prop(hb3, "strength_adjust")

    @staticmethod
    def draw_common_cursor(context, xy, tool):
        global G_on_mesh
        global G_loc_world
        global G_nor_world
        global G_rad_3d
        global G_brush_size
        global G_source_circle, G_show_clone_circle

        brush = context.tool_settings.particle_edit.brush
        region = context.region
        rv3d = context.region_data
        if not hasattr(context.scene, "hair_brush_3d"):
            return

        hb3 = context.scene.hair_brush_3d
        if tool.support_3d and tool.can_update_3d_cursor(hb3):
            on_mesh, loc, nor = isect_object_from_2d_coord(region, rv3d, context.object, xy)
            if on_mesh:
                G_on_mesh = True
                obmat = context.object.matrix_world
                G_loc_world = obmat @ loc
                G_nor_world = obmat.to_quaternion() @ nor
            else:
                G_on_mesh = False
                G_loc_world.zero()
                G_nor_world.zero()

            if hb3.use_3d_brush:
                if on_mesh:
                    clone_source = None
                    if not hb3.scale_size or brush.size != G_brush_size:
                        G_rad_3d, hb3.unprojected_rad = calc_unprojected_radius(obmat, loc, region, rv3d, brush.size)
                        if brush.size != G_brush_size:
                            G_brush_size = brush.size
                    else:
                        brush.size = G_brush_size = calc_projected_radius(obmat, region, rv3d, hb3.unprojected_rad,
                                                                          G_loc_world)
                    if G_show_clone_circle and G_source_circle and hb3.clone_realtime_show:
                        clone_source = particle_hook.get_clone_source_for_draw(
                            context.scene.as_pointer(),
                            context.object.as_pointer(),
                            context.object.particle_systems.active.as_pointer(),
                            G_loc_world.to_tuple(),
                            G_nor_world.to_tuple(),
                            G_rad_3d,
                            get_enum_item_num(CloneCoordinateEnumItems, hb3.clone_coordinate, 0),
                            hb3.clone_rotation_angle)
                    draw_circle_cursor_3d(context, obmat, loc, nor, hb3.unprojected_rad, clone_source)
                    return
                else:
                    if not hb3.auto_switch_mode:
                        return

        draw_circle_cursor_2d(xy, brush.size)

    def draw_cursor(context, tool, xy):
        if HairBrush3D.poll(context.object) and BrushType.contains(tool.idname):
            HairBrush3D.draw_common_cursor(context, xy, BrushType.get(tool.idname))


class BaseOperator(Operator):
    bl_label = ""
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context: bpy.types.Context):
        obj = context.object
        return obj and obj.mode == 'PARTICLE_EDIT' and obj.particle_systems.active.settings.type == 'HAIR'


class PARTICLE_OT_HairBrush3DEdit(BaseOperator):
    bl_idname = "particle.brush3d_edit"
    bl_label = "Brush 3D Edit"
    bl_description = "Apply brush 3D edit"

    def get_view3d_clip_range(self, context: bpy.types.Context):
        clip_start = 0.1
        clip_end = 100

        v3d = context.space_data
        rv3d = context.region_data
        if v3d and rv3d:
            clip_start, clip_end = v3d.clip_start, v3d.clip_end
            if rv3d.view_perspective == 'CAMERA':
                if v3d.camera and v3d.camera.type == 'CAMERA':
                    deps = context.evaluated_depsgraph_get()
                    cam = v3d.camera.evaluated_get(deps)
                    clip_start, clip_end = cam.data.clip_start, cam.data.clip_end
            elif rv3d.view_perspective == 'ORTHO':
                clip_end *= 0.5
                clip_start = -clip_end
        return clip_start, clip_end

    def create_mesh_surface_batch(self, mesh):
        mesh.calc_loop_triangles()

        vertices = np.empty((len(mesh.vertices), 3), 'f')
        indices = np.empty((len(mesh.loop_triangles), 3), 'i')

        mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3))
        mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3))

        fmt = gpu.types.GPUVertFormat()
        fmt.attr_add(id="pos", comp_type='F32', len=3, fetch_mode='FLOAT')
        vbo = gpu.types.GPUVertBuf(len=len(vertices), format=fmt)
        vbo.attr_fill("pos", vertices)
        ibo = gpu.types.GPUIndexBuf(type='TRIS', seq=indices)
        return gpu.types.GPUBatch(type='TRIS', buf=vbo, elem=ibo)

    def get_object_depth_buffer(self, context, obj):
        region = context.region
        rv3d = context.region_data

        size = (region.width, region.height)

        bs = gpu.shader.from_builtin('2D_UNIFORM_COLOR')
        batch = batch_for_shader(bs, 'POINTS', {"pos": ((1, 1),)})
        shader = gpu.types.GPUShader(vertex_shader, fragment_shader)
        depth_tex = gpu.types.GPUTexture(size, format='DEPTH_COMPONENT24')
        frame_buf = gpu.types.GPUFrameBuffer(depth_slot=depth_tex)
        # frame_buf.viewport_set(region.x, region.y, size[0], size[1])
        depth_buf = None

        with frame_buf.bind():
            gpu.state.depth_test_set('LESS_EQUAL')
            frame_buf.clear(color=(0, 0, 0, 1), depth=1.0)
            gpu.state.depth_mask_set(True)
            # gpu.state.color_mask_set(False, False, False, False)
            with gpu.matrix.push_pop():
                gpu.matrix.load_identity()
                gpu.matrix.load_projection_matrix(rv3d.window_matrix)
                gpu.matrix.load_matrix(rv3d.view_matrix)
                gpu.matrix.multiply_matrix(obj.matrix_world)
                ori_batch = particle_hook.set_mesh_surface_gpu_batch(batch, obj.data.as_pointer())
                if ori_batch:
                    batch.draw(shader)
                else:
                    mesh_batch = self.create_mesh_surface_batch(obj.data)
                    mesh_batch.draw(shader)
                gpu.matrix.load_matrix(rv3d.view_matrix)
                if ori_batch:
                    particle_hook.restore_gpu_batch(batch, ori_batch)
                depth_buf = depth_tex.read()
            gpu.state.depth_mask_set(False)
            gpu.state.depth_test_set('NONE')

        return depth_buf

    def ensure_apply_do(self):
        if not self.brush_support_3d:
            return True
        if self.use_3d:
            return not (self.start_on_mesh ^ G_on_mesh)
        return self.brush_support_2d

    def brush_edit_exit(self, context):
        if self.edit_data:
            particle_hook.free_brush_edit_data(self.edit_data)
            self.edit_data = None
        if self.cur_active_tool == "hair_brush_3d.Color":
            ColorGroupManager.generate_color_groups(context.object.particle_systems.active)
            update_hair_color(context, ColorUpdateFlag.FORCE_TEX)

    def apply_brush_edit(self, event: bpy.types.Event):
        if particle_hook.apply_brush_edit(
                self.edit_data,
                (event.mouse_region_x, event.mouse_region_y),
                G_loc_world.to_tuple(),
                G_nor_world.to_tuple(),
                event.shift):
            self.obj.update_tag(refresh={'OBJECT', 'DATA'})
            return True
        return False

    def modal(self, context, event: bpy.types.Event):
        context.area.tag_redraw()
        if event.type in ('LEFTMOUSE', 'MIDDLEMOUSE', 'RIGHTMOUSE'):
            if event.value == 'RELEASE':
                self.brush_edit_exit(context)
                return {'FINISHED'}
        elif event.type == 'MOUSEMOVE' and self.ensure_apply_do():
            if (self.apply_brush_edit(event)
                    and self.cur_active_tool == "hair_brush_3d.Color"
                    and self.update_color_realtime):
                update_hair_color(context, ColorUpdateFlag.FORCE_TEX)
        return {'RUNNING_MODAL'}

    def build_edit_args(self, context, hb3, brush_type, flag):
        edit_args = dict()
        edit_args['brush_type'] = brush_type.brush_type
        edit_args['rad_3d'] = G_rad_3d
        edit_args['flag'] = flag
        edit_args['strength_adjust'] = hb3.strength_adjust
        edit_args['smooth_mode'] = 0 if hb3.smooth_mode == 'SHAPE' else 1
        edit_args['limited_length'] = hb3.limited_length
        edit_args['orient_mode'] = get_enum_item_num(OrientModeEnumItems, hb3.orient_mode, 0)
        edit_args['curving'] = hb3.curving
        edit_args['weight_mode'] = get_enum_item_num(WeightModeEnumItems, hb3.weight_mode, 0)
        edit_args['attract_method'] = 0 if hb3.attract_method == "HEAD" else 1
        edit_args['continuous_attract'] = hb3.continuous_attract
        edit_args['random_control'] = hb3.random_control
        edit_args['percent_control'] = hb3.percent_control
        edit_args['attract_mode'] = 0 if hb3.attract_mode == "POINT" else 1
        edit_args['head_strength'] = hb3.head_strength
        edit_args['middle_strength'] = hb3.middle_strength
        edit_args['rough_seed'] = hb3.rough_seed
        edit_args['noise_flag'] = get_noise_flag(hb3)
        edit_args['clip_range'] = self.get_view3d_clip_range(context)
        edit_args['clone_mode'] = 0 if hb3.clone_mode == "CLONE_DRAW" else 1
        edit_args['clone_coordinate'] = get_enum_item_num(CloneCoordinateEnumItems, hb3.clone_coordinate, 0)
        edit_args['clone_rotation_angle'] = hb3.clone_rotation_angle
        edit_args['brush_color'] = (hb3.brush_color.r, hb3.brush_color.g, hb3.brush_color.b)
        edit_args['falloff_shape'] = 0 if hb3.falloff_shape == "SPHERE" else 1
        brush_data = HairBrush3DCommonProps.get_brush_data(hb3)
        if hb3.brush_data is None:
            hb3.brush_data = brush_data
        brush_data.curve.update()
        edit_args['brush_data'] = brush_data.as_pointer()
        return edit_args

    def invoke(self, context: bpy.types.Context, event):
        if not HairBrush3D.poll(context.object):
            return {'CANCELLED'}

        hb3 = context.scene.hair_brush_3d
        self.cur_active_tool = ToolSelectPanelHelper.tool_active_from_context(context).idname
        if not BrushType.contains(self.cur_active_tool):
            return {'CANCELLED'}

        brush_type = BrushType.get(self.cur_active_tool)
        if brush_type.can_cancel(hb3):
            return {'CANCELLED'}

        if (self.cur_active_tool == "hair_brush_3d.Clone" and hb3.clone_mode != "CLONE_SOURCE"
                and not event.shift and not G_source_circle):
            return {'CANCELLED'}

        pe = context.tool_settings.particle_edit
        if brush_type.builtin_brush_id or pe.tool == 'WEIGHT':
            brush = pe.brush
            cur_count = brush.count
            cur_size = brush.size
            cur_steps = brush.steps
            cur_strength = brush.strength
            bpy.ops.wm.tool_set_by_id(name=brush_type.builtin_brush_id if brush_type.
                                      builtin_brush_id else "builtin_brush.Comb")
            brush = pe.brush
            brush.count = cur_count
            brush.size = cur_size
            brush.steps = cur_steps
            brush.strength = cur_strength
            if brush_type.can_to_builtin(hb3):
                bpy.ops.particle.brush_edit('INVOKE_DEFAULT')
                bpy.ops.wm.tool_set_by_id(name=self.cur_active_tool)
                return {'FINISHED'}
            bpy.ops.wm.tool_set_by_id(name=self.cur_active_tool)

        self.obj = context.object
        deps = context.evaluated_depsgraph_get()
        evaled_obj = context.object.evaluated_get(depsgraph=deps)
        depth_buf = self.get_object_depth_buffer(context, evaled_obj)
        if depth_buf is None:
            self.report({'ERROR'}, "Failed to get the depth of the object, please make sure it's visible in view")
            return {'CANCELLED'}

        psys = context.object.particle_systems.active
        par_num = len(psys.particles)
        if self.cur_active_tool == "hair_brush_3d.Color" and par_num:
            if "hair_color__" not in psys.settings:
                psys.settings["hair_color__"] = [0.0] * (3 * par_num)
            elif len(psys.settings["hair_color__"]) != 3 * par_num:
                self.report({'ERROR'},
                            "The hair count is different from that recorded in the color groups. Please delete all color groups and then repaint the hair.")
                return {"CANCELLED"}

        self.brush_support_2d = brush_type.support_2d
        self.brush_support_3d = brush_type.support_3d
        self.start_on_mesh = False
        self.use_3d = False
        flag = 0
        if (hb3.use_3d_brush and self.brush_support_3d) or not self.brush_support_2d:
            flag |= EditArgsFlag.USE_3D
            self.use_3d = True
        if G_on_mesh:
            self.start_on_mesh = True
            if self.brush_support_3d:
                flag |= EditArgsFlag.ON_MESH
        if hb3.scale_size:
            flag |= EditArgsFlag.USE_SCALE

        edit_args = self.build_edit_args(context, hb3, brush_type, flag)
        if self.cur_active_tool == "hair_brush_3d.SnakeHook" and hb3.falloff_shape == 'SPHERE':
            on_mesh, loc, nor = isect_object_from_2d_coord(
                context.region, context.region_data, context.object, (event.mouse_region_x, event.mouse_region_y))
            if on_mesh:
                global G_loc_world, G_nor_world
                edit_args['flag'] |= EditArgsFlag.ON_MESH
                obmat = context.object.matrix_world
                G_loc_world = obmat @ loc
                G_nor_world = obmat.to_quaternion() @ nor
            edit_args['widget_unit'] = particle_hook.get_widget_unit(context.preferences.as_pointer())

        self.edit_data = particle_hook.init_brush_edit_data(
            deps.as_pointer(),
            context.scene.as_pointer(),
            self.obj.as_pointer(),
            psys.as_pointer(),
            context.region.as_pointer(),
            context.space_data.as_pointer(),
            depth_buf,
            edit_args
        )
        if self.edit_data is None:
            return {'CANCELLED'}

        if self.cur_active_tool == "hair_brush_3d.Clone":
            if hb3.clone_mode == "CLONE_SOURCE" or event.shift:
                global G_pack_source_circle, G_show_clone_circle
                G_pack_source_circle = True
                G_show_clone_circle = particle_hook.set_clone_source(self.edit_data,
                                                                     (event.mouse_region_x, event.mouse_region_y),
                                                                     G_loc_world.to_tuple(),
                                                                     G_nor_world.to_tuple())
            else:
                self.apply_brush_edit(event)
                context.area.tag_redraw()
            self.brush_edit_exit(context)
            return {'FINISHED'}

        if self.cur_active_tool == "hair_brush_3d.Color":
            ColorGroupManager.backup_hair_color(psys.settings, backup_all=True)
        self.update_color_realtime = hb3.update_color_realtime
        self.apply_brush_edit(event)

        context.window_manager.modal_handler_add(self)
        return {'RUNNING_MODAL'}


class PARTICLE_OT_HairBrush3DCloneRotate(bpy.types.Operator):
    bl_idname = "particle.brush3d_clone_rotate"
    bl_label = "Rotate Clone Brush"
    bl_description = "Rotate clone brush"
    bl_options = {"REGISTER", "UNDO"}

    rotate_mode: EnumProperty(
        name="Rotate mode",
        items=[("ADD", "Add", "Add rotation", 0),
               ("SUB", "Sub", "Sub rotation", 1)],
        options={'HIDDEN'})

    @classmethod
    def poll(cls, context: bpy.types.Context):
        if not HairBrush3D.poll(context.object):
            return False
        tc = ToolSelectPanelHelper.tool_active_from_context(context)
        return tc.idname == "hair_brush_3d.Clone"

    def execute(self, context):
        hb3 = context.scene.hair_brush_3d
        if self.rotate_mode == 'ADD':
            hb3.clone_rotation_angle += 1
        else:
            hb3.clone_rotation_angle -= 1
        return {'FINISHED'}


class PARTICLE_TOOL_Brush3DComb(HairBrush3D):
    # The prefix of the idname should be your add-on name.
    bl_idname = "hair_brush_3d.Comb"
    bl_label = "Comb"
    bl_icon = "brush3d.particle.comb"

    builtin_brush_id = "builtin_brush.Comb"

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout)
        settings = context.tool_settings.particle_edit
        layout.prop(settings, "use_emitter_deflect", text="Deflect Emitter")
        sub = layout.row()
        sub.active = settings.use_emitter_deflect
        sub.prop(settings, "emitter_distance", text="Distance")


class PARTICLE_TOOL_Brush3DSmooth(HairBrush3D):
    bl_idname = "hair_brush_3d.Smooth"
    bl_label = "Smooth"
    bl_icon = "brush3d.particle.smooth"

    builtin_brush_id = "builtin_brush.Smooth"

    @classmethod
    def can_to_builtin(cls, hb3):
        return hb3.smooth_mode == 'SHAPE' and HairBrush3D.can_to_builtin(hb3)

    def draw_settings(context, layout, tool):
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)
        hb3 = context.scene.hair_brush_3d
        layout.prop(hb3, "smooth_mode", expand=True)


class PARTICLE_TOOL_Brush3DAdd(HairBrush3D):
    bl_idname = "hair_brush_3d.Add"
    bl_label = "Add"
    bl_icon = "brush3d.particle.add"

    builtin_brush_id = "builtin_brush.Add"

    @classmethod
    def can_cancel(cls, hb3):
        return hb3.use_3d_brush and not G_on_mesh

    def draw_settings(context, layout, tool):
        HairBrush3D.draw_common_options(context, layout, False, draw_adjust=False)
        settings = context.tool_settings.particle_edit
        brush = settings.brush
        layout.prop(brush, "count")
        layout.prop(settings, "use_default_interpolate")
        layout.prop(brush, "steps", slider=True)
        layout.prop(settings, "default_key_count", slider=True)


class PARTICLE_TOOL_Brush3DLength(HairBrush3D):
    bl_idname = "hair_brush_3d.Length"
    bl_label = "Length"
    bl_description = "Works in 2D mode only"
    bl_icon = "brush3d.particle.length"

    builtin_brush_id = "builtin_brush.Length"
    support_3d = False

    @classmethod
    def can_cancel(cls, hb3):
        return False

    @classmethod
    def can_to_builtin(cls, hb3):
        return hb3.limited_length == 0.0

    def draw_settings(context, layout, tool):
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)
        brush = context.tool_settings.particle_edit.brush
        hb3 = context.scene.hair_brush_3d
        layout.row().prop(brush, "length_mode", expand=True)
        layout.prop(hb3, "limited_length")

    def draw_cursor(context, tool, xy):
        if HairBrush3D.poll(context.object):
            draw_circle_cursor_2d(xy, context.tool_settings.particle_edit.brush.size)


class PARTICLE_TOOL_Brush3DPuff(HairBrush3D):
    bl_idname = "hair_brush_3d.Puff"
    bl_label = "Puff"
    bl_icon = "brush3d.particle.puff"

    builtin_brush_id = "builtin_brush.Puff"

    def draw_settings(context, layout, tool):
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)
        brush = context.tool_settings.particle_edit.brush
        layout.row().prop(brush, "puff_mode", expand=True)
        layout.prop(brush, "use_puff_volume")


class PARTICLE_TOOL_Brush3DSnakeHook(HairBrush3D):
    bl_idname = "hair_brush_3d.SnakeHook"
    bl_label = "Snake Hook"
    bl_icon = "brush3d.particle.snake_hook"

    support_3d = False

    @classmethod
    def can_cancel(cls, hb3):
        return False

    @classmethod
    def draw_falloff_settings(cls, layout, brush, prop_text=""):
        col = layout.column(align=True)
        row = col.row(align=True)
        row.prop(brush, "curve_preset", text=prop_text)

        if brush.curve_preset == 'CUSTOM':
            layout.template_curve_mapping(brush, "curve", brush=True)

            col = layout.column(align=True)
            row = col.row(align=True)
            row.operator(PARTICLE_OT_BrushCurvePreset.bl_idname, icon='SMOOTHCURVE', text="").shape = 'SMOOTH'
            row.operator(PARTICLE_OT_BrushCurvePreset.bl_idname, icon='SPHERECURVE', text="").shape = 'ROUND'
            row.operator(PARTICLE_OT_BrushCurvePreset.bl_idname, icon='ROOTCURVE', text="").shape = 'ROOT'
            row.operator(PARTICLE_OT_BrushCurvePreset.bl_idname, icon='SHARPCURVE', text="").shape = 'SHARP'
            row.operator(PARTICLE_OT_BrushCurvePreset.bl_idname, icon='LINCURVE', text="").shape = 'LINE'
            row.operator(PARTICLE_OT_BrushCurvePreset.bl_idname, icon='NOCURVE', text="").shape = 'MAX'

    def draw_settings(context, layout, tool):
        region_type = context.region.type
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)
        hb3 = context.scene.hair_brush_3d
        layout.prop(hb3, "falloff_shape", expand=True)
        if region_type == 'TOOL_HEADER':
            layout.popover("VIEW3D_PT_tools_particle_brush_falloff")
        else:
            PARTICLE_TOOL_Brush3DSnakeHook.draw_falloff_settings(
                layout, HairBrush3DCommonProps.get_brush_data(hb3), "Falloff Curve")

    def draw_cursor(context, tool, xy):
        if HairBrush3D.poll(context.object):
            draw_circle_cursor_2d(xy, context.tool_settings.particle_edit.brush.size)


class PARTICLE_TOOL_Brush3DOrient(HairBrush3D):
    bl_idname = "hair_brush_3d.Orient"
    bl_label = "Orient"
    bl_icon = "brush3d.particle.orient"

    @classmethod
    def can_update_3d_cursor(cls, hb3):
        return hb3.use_3d_brush or hb3.orient_mode in {'SUNDER', 'REPEL'}

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout)
        hb3 = context.scene.hair_brush_3d
        layout.prop(hb3, "orient_mode", expand=True)
        layout.prop(hb3, "curving")


class PARTICLE_TOOL_Brush3DBend(HairBrush3D):
    bl_idname = "hair_brush_3d.Bend"
    bl_label = "Bend"
    bl_icon = "brush3d.particle.bend"

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)


class PARTICLE_TOOL_Brush3DPart(HairBrush3D):
    bl_idname = "hair_brush_3d.Part"
    bl_label = "Part"
    bl_icon = "brush3d.particle.part"

    @classmethod
    def can_update_3d_cursor(cls, hb3):
        return True

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)


class PARTICLE_TOOL_Brush3DWeight(HairBrush3D):
    bl_idname = "hair_brush_3d.Weight"
    bl_label = "Weight"
    bl_description = "Works in 2D mode only"
    bl_icon = "brush3d.particle.weight"

    builtin_brush_id = "builtin_brush.Weight"
    support_3d = False

    @classmethod
    def can_cancel(cls, hb3):
        return False

    @classmethod
    def can_to_builtin(cls, hb3):
        return hb3.weight_mode == 'DEFAULT'

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout, draw_adjust=False)
        hb3 = context.scene.hair_brush_3d
        layout.prop(hb3, "weight_mode", expand=True)

    def draw_cursor(context, tool, xy):
        if HairBrush3D.poll(context.object):
            draw_circle_cursor_2d(xy, context.tool_settings.particle_edit.brush.size)


class PARTICLE_TOOL_Brush3DNoise(HairBrush3D):
    bl_idname = "hair_brush_3d.Noise"
    bl_label = "Noise"
    bl_icon = "brush3d.particle.noise"

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        hb3 = context.scene.hair_brush_3d
        if not hb3.use_rough:
            HairBrush3D.draw_common_options(context, layout, True, False)
            layout.prop(hb3, "use_rough")
            layout.prop(hb3, "head_strength", slider=True)
            layout.prop(hb3, "middle_strength", slider=True)
        else:
            HairBrush3D.draw_common_options(context, layout, False, False)
            brush = context.tool_settings.particle_edit.brush
            layout.prop(hb3, "use_rough")
            layout.prop(hb3, "reset_pose")
            layout.prop(hb3, "head_strength", text="Uniform", slider=True)
            layout.prop(hb3, "middle_strength", text="Endpoint", slider=True)
            layout.prop(brush, "strength", text="Random", slider=True)
            if not hb3.automatic_seed:
                layout.prop(hb3, "rough_seed")
            layout.prop(hb3, "automatic_seed")

        row = layout.row(align=True)
        row.label(icon='LOCKED')
        sub = row.row(align=True)
        sub.scale_x = 0.7
        sub.prop(hb3, "use_lock_x", text="X", toggle=True)
        sub.prop(hb3, "use_lock_y", text="Y", toggle=True)
        sub.prop(hb3, "use_lock_z", text="Z", toggle=True)


class PARTICLE_TOOL_Brush3DAttract(HairBrush3D):
    bl_idname = "hair_brush_3d.Attract"
    bl_label = "Attract"
    bl_icon = "brush3d.particle.attract"

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout)
        hb3 = context.scene.hair_brush_3d
        row = layout.row(align=True)
        row.scale_x = 0.8
        row.prop(hb3, "attract_method")
        layout.prop(hb3, "continuous_attract")
        layout.prop(hb3, "random_control", slider=True)
        layout.prop(hb3, "percent_control", slider=True)
        layout.prop(hb3, "attract_mode", expand=True)


class PARTICLE_TOOL_Brush3DClone(HairBrush3D):
    bl_idname = "hair_brush_3d.Clone"
    bl_label = "Clone"
    bl_icon = "brush3d.particle.clone"
    bl_description = "Works in 3D mode only"
    bl_keymap = (
        ("particle.brush3d_edit", {"type": 'LEFTMOUSE', "value": 'PRESS'}, {"properties": []}),
        ("particle.brush3d_edit", {"type": 'LEFTMOUSE', "value": 'PRESS', "shift": True}, {"properties": []}),
        ("particle.brush3d_clone_rotate",
         {"type": 'WHEELUPMOUSE', "value": 'PRESS', "ctrl": True},
         {"properties": [("rotate_mode", 'ADD')]}),
        ("particle.brush3d_clone_rotate",
         {"type": 'WHEELDOWNMOUSE', "value": 'PRESS', "ctrl": True},
         {"properties": [("rotate_mode", 'SUB')]}),)

    support_2d = False

    @classmethod
    def can_cancel(cls, hb3):
        return not hb3.use_3d_brush or not G_on_mesh

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout, False, False)
        hb3 = context.scene.hair_brush_3d
        layout.prop(hb3, "clone_mode", expand=True)
        row = layout.row(align=True)
        row.scale_x = 0.85
        row.prop(hb3, "clone_coordinate")
        layout.prop(hb3, "clone_rotation_angle")
        layout.prop(hb3, "clone_realtime_show")


class PARTICLE_TOOL_Brush3DColor(HairBrush3D):
    bl_idname = "hair_brush_3d.Color"
    bl_label = "Color"
    bl_description = "Works in 2D mode only"
    bl_icon = "brush3d.particle.color"

    support_3d = False

    @classmethod
    def can_cancel(cls, hb3):
        return False

    def draw_settings(context: bpy.types.Context, layout: bpy.types.UILayout, tool):
        HairBrush3D.draw_common_options(context, layout, False, False)
        hb3 = context.scene.hair_brush_3d
        layout.prop(hb3, "show_hair_color")
        layout.prop(hb3, "update_color_realtime")
        row = layout.row(align=True)
        row.scale_x = 0.58
        row.prop(hb3, "brush_color")

    def draw_cursor(context, tool, xy):
        if HairBrush3D.poll(context.object):
            draw_circle_cursor_2d(xy, context.tool_settings.particle_edit.brush.size)


class VIEW3D_PT_tools_hair_brush_3d_options(Panel):
    bl_label = "3D Brush Options"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_context = ".particlemode"
    bl_category = "Tool"
    bl_parent_id = "VIEW3D_PT_tools_particlemode_options"

    @classmethod
    def poll(cls, context):
        tc = ToolSelectPanelHelper.tool_active_from_context(context)
        return BrushType.contains(tc.idname) and BrushType.get(tc.idname).support_3d

    def draw(self, context):
        tc = ToolSelectPanelHelper.tool_active_from_context(context)
        if tc.idname == "hair_brush_3d.Color":
            return
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False  # No animation.
        hb3 = context.scene.hair_brush_3d

        layout.prop(hb3, "use_3d_brush")
        col = layout.column()
        col.prop(hb3, "auto_switch_mode")
        col.prop(hb3, "scale_size")
        col.enabled = hb3.use_3d_brush


class PARTICLE_OT_DeleteColorGroup(BaseOperator):
    bl_idname = "particle.delete_color_group"
    bl_description = "Delete the active color group and reset the hair color"

    index: IntProperty(default=0, options={'HIDDEN', 'SKIP_SAVE'})

    def execute(self, context: bpy.types.Context):
        psys = context.object.particle_systems.active
        pset = psys.settings
        groups = pset.hair_color_groups
        if self.index >= len(groups):
            return {'CANCELLED'}

        if self.index == -1 or len(groups) <= 1:
            ColorGroupManager.backup_hair_color(pset, backup_all=True)
            ColorGroupManager.clear_color_groups(pset, ColorUpdateFlag.FORCE_TEX, context)
        else:  # self.index < len(groups)
            rgb = tuple(groups[self.index].rgb)
            ColorGroupManager.backup_hair_color(pset, rgb=rgb)
            ColorGroupManager.modify_color_groups(self, context, rgb, HairColorFlag.OP_DELETE, psys,
                                                  ColorUpdateFlag.FORCE_TEX)
        context.area.tag_redraw()
        return {'FINISHED'}


class PARTICLE_OT_UndoDelColorGroup(BaseOperator):
    bl_label = "Undo"
    bl_idname = "particle.undo_hair_color"
    bl_description = "Undo painting hair color or deleting color group. Note, once you change the amount of hairs then you can't undo anymore"

    @classmethod
    def poll(cls, context: bpy.types.Context):
        return BaseOperator.poll(context) and ColorGroupManager.undo_stack

    def execute(self, context: bpy.types.Context):
        psys = context.object.particle_systems.active
        if len(psys.particles) * 3 != ColorGroupManager.undo_stack[-1][-1]:
            ColorGroupManager.undo_stack = []
            self.report({'ERROR'}, "The count of the hair has changed, can't undo!")
            return {'CANCELLED'}

        ColorGroupManager.undo_hair_color(psys, context)
        context.area.tag_redraw()
        return {'FINISHED'}


class PARTICLE_OT_AssignColor2Selected(BaseOperator):
    bl_label = "Assign Color"
    bl_idname = "particle.assign_color2selected"
    bl_description = "Assign brush color to the selected hair"

    show_ui: BoolProperty(default=False, options={'HIDDEN', 'SKIP_SAVE'})

    def invoke(self, context: bpy.types.Context, event):
        if self.show_ui:
            return context.window_manager.invoke_props_popup(self, event)
        else:
            return self.execute(context)

    def draw(self, context):
        if self.show_ui:
            hb3 = context.scene.hair_brush_3d
            self.layout.prop(hb3, "brush_color")

    def execute(self, context: bpy.types.Context):
        psys = context.object.particle_systems.active
        pset = psys.settings
        if 'hair_color__' not in pset:
            pset["hair_color__"] = [0.0] * (3 * len(psys.particles))
        rgb = tuple(context.scene.hair_brush_3d.brush_color)
        ColorGroupManager.modify_color_groups(self, context, rgb, HairColorFlag.OP_ASSIGN, psys,
                                              ColorUpdateFlag.FORCE_TEX)

        context.area.tag_redraw()
        return {'FINISHED'}


class PARTICLE_OT_SelectInGroup(BaseOperator):
    bl_idname = "particle.select_hair_in_group"
    bl_label = ""
    option: IntProperty(default=HairColorFlag.OP_SELECT_ONLY, options={'HIDDEN', 'SKIP_SAVE'})

    @classmethod
    def description(cls, context, props):
        if props.option == HairColorFlag.OP_SELECT_ONLY:
            return "Select the hair in the active color group"
        else:
            return "Add the hair in the active color group to selection"

    def execute(self, context: bpy.types.Context):
        psys = context.object.particle_systems.active
        groups = psys.settings.hair_color_groups
        idx = psys.settings.color_group_idx
        if 0 <= idx < len(groups):
            rgb = tuple(groups[idx].rgb)
            particle_hook.set_hair_color_group(psys.as_pointer(), rgb, self.option)
            select_mode = context.tool_settings.particle_edit.select_mode
            context.tool_settings.particle_edit.select_mode = select_mode
        return {'FINISHED'}


class PARTICLE_OT_BrushCurvePreset(BaseOperator):
    """Set brush shape"""
    bl_idname = "particle.brush_curve_preset"
    bl_label = "Preset"
    bl_options = set()

    shape: EnumProperty(
        name="Mode",
        items=[('SMOOTH', "SMOOTH", "", 0),
               ('ROUND', "Round", "", 1),
               ('ROOT', "Root", "", 2),
               ('SHARP', "Sharp", "", 3),
               ('LINE', "Line", "", 4),
               ('MAX', "Max", "", 5)])

    def ensure_points_num(self, curve: bpy.types.CurveMap, num):
        while len(curve.points) < num:
            curve.points.new(1, 1)
        while len(curve.points) > num:
            curve.points.remove(curve.points[-1])

    def execute(self, context: bpy.types.Context):
        hb3 = context.scene.hair_brush_3d
        brush = HairBrush3DCommonProps.get_brush_data(hb3)
        if hb3.brush_data is None:
            hb3.brush_data = brush
        if len(brush.curve.curves) == 0:
            brush.curve.initialize()
        curve = brush.curve.curves[0]
        if self.shape == 'SMOOTH':
            self.ensure_points_num(curve, 4)
            curve.points[0].location = 0, 1
            curve.points[1].location = 0.25, 0.94
            curve.points[2].location = 0.75, 0.06
            curve.points[3].location = 1, 0
        elif self.shape == 'ROUND':
            self.ensure_points_num(curve, 4)
            curve.points[0].location = 0, 1
            curve.points[1].location = 0.5, 0.9
            curve.points[2].location = 0.86, 0.5
            curve.points[3].location = 1, 0
        elif self.shape == 'ROOT':
            self.ensure_points_num(curve, 4)
            curve.points[0].location = 0, 1
            curve.points[1].location = 0.25, 0.95
            curve.points[2].location = 0.75, 0.44
            curve.points[3].location = 1, 0
        elif self.shape == 'SHARP':
            self.ensure_points_num(curve, 4)
            curve.points[0].location = 0, 1
            curve.points[1].location = 0.25, 0.5
            curve.points[2].location = 0.75, 0.04
            curve.points[3].location = 1, 0
        elif self.shape == 'LINE':
            self.ensure_points_num(curve, 2)
            curve.points[0].location = brush.curve.clip_min_x, brush.curve.clip_max_y
            curve.points[1].location = brush.curve.clip_max_x, brush.curve.clip_min_y
        elif self.shape == 'MAX':
            self.ensure_points_num(curve, 2)
            curve.points[0].location = 0, 1
            curve.points[1].location = 1, 1
        brush.curve.update()
        return {'FINISHED'}


class PARTICLE_UL_color_group(bpy.types.UIList):
    def draw_item(self, context, layout, data, item, icon, active_data, active_propname):
        row = layout.row(align=True)
        row.label(text=item.name)
        row.prop(item, "rgb_show", text="")


class VIEW3D_PT_hair_color_options(Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "Tool"
    bl_context = ".particlemode"
    bl_label = "Hair Color Groups"

    # bl_options = {'DEFAULT_CLOSED'}

    @classmethod
    def poll(cls, context: bpy.types.Context):
        return context.object.particle_systems.active.settings.type == 'HAIR'

    def draw(self, context):
        layout = self.layout
        layout.use_property_decorate = False  # No animation.
        self.draw_layout(layout, context)

    @classmethod
    def draw_layout(cls, layout, context, show_in_piemenu=False):
        hb3 = context.scene.hair_brush_3d
        pset = context.object.particle_systems.active.settings
        layout.prop(hb3, "show_hair_color", emboss=not show_in_piemenu)
        if not show_in_piemenu:
            layout.separator()
        row = layout.row(align=True)
        row.prop(hb3, "brush_color", text="")
        if show_in_piemenu:
            row.operator(PARTICLE_OT_AssignColor2Selected.bl_idname, text="Change Color").show_ui = True
            layout.separator()
            layout.operator(PARTICLE_OT_AssignColor2Selected.bl_idname).show_ui = False
            layout.separator()
        else:
            row.operator(PARTICLE_OT_AssignColor2Selected.bl_idname).show_ui = False
        layout.template_list("PARTICLE_UL_color_group", "", pset,
                             "hair_color_groups", pset, "color_group_idx")
        if "hair_color__" in pset and pset["hair_color__"]:
            row = layout.row(align=True)
            row.operator(PARTICLE_OT_SelectInGroup.bl_idname, text="Select  ").option = HairColorFlag.OP_SELECT_ONLY
            row.operator(PARTICLE_OT_SelectInGroup.bl_idname, text="Add  ").option = HairColorFlag.OP_SELECT_ADD
            row = layout.row(align=True)
            row.operator(PARTICLE_OT_DeleteColorGroup.bl_idname, text="Delete All").index = -1
            row.operator(PARTICLE_OT_DeleteColorGroup.bl_idname, text="Delete Active").index = pset.color_group_idx
            row.operator(PARTICLE_OT_UndoDelColorGroup.bl_idname)
        elif ColorGroupManager.undo_stack:
            layout.operator(PARTICLE_OT_UndoDelColorGroup.bl_idname)


class VIEW3D_PT_tools_particle_brush_falloff(Panel):
    bl_context = ".particlemode"  # dot on purpose (access from topbar)
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOL_HEADER'
    bl_label = "Falloff"
    bl_options = {'DEFAULT_CLOSED'}

    def draw(self, context):
        PARTICLE_TOOL_Brush3DSnakeHook.draw_falloff_settings(
            self.layout,
            HairBrush3DCommonProps.get_brush_data(context.scene.hair_brush_3d))


def on_object_mode_change():
    context = bpy.context
    obj = context.object
    if not obj or obj.mode != 'PARTICLE_EDIT':
        return
    psys = obj.particle_systems.active
    ColorGroupManager.undo_stack = []
    if psys.settings.type != 'HAIR':
        HairRootsTree.reset_roots_tree()
        return
    if not psys.is_edited:
        ColorGroupManager.clear_color_groups(psys.settings, ColorUpdateFlag.FORCE_TEX | ColorUpdateFlag.FROM_DEPS, None)

    HairRootsTree.generate_hair_roots(psys.as_pointer(), obj.evaluated_get(context.evaluated_depsgraph_get()))


def on_particle_edit_tool_change():
    context = bpy.context
    obj = context.object
    if not obj or obj.mode != 'PARTICLE_EDIT':
        return
    psys = obj.particle_systems.active
    if psys.settings.type != 'HAIR':
        return
    bpy.ops.wm.tool_set_by_id(name="builtin_brush." +
                                   context.tool_settings.particle_edit.tool.title(),
                              space_type='VIEW_3D')


def register_icons():
    icons_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "icons")
    if not os.path.exists(icons_dir):
        return

    for fn in os.listdir(icons_dir):
        icon_name = os.path.splitext(fn)[0]
        if icon_name in _icon_cache:
            continue

        try:
            icon_value = bpy.app.icons.new_triangles_from_file(os.path.join(icons_dir, fn))
            _icon_cache[icon_name] = icon_value
        except Exception as ex:
            pass


tools = (PARTICLE_TOOL_Brush3DComb, PARTICLE_TOOL_Brush3DSmooth,
         PARTICLE_TOOL_Brush3DAdd, PARTICLE_TOOL_Brush3DLength,
         PARTICLE_TOOL_Brush3DPuff, PARTICLE_TOOL_Brush3DSnakeHook,
         PARTICLE_TOOL_Brush3DOrient, PARTICLE_TOOL_Brush3DBend,
         PARTICLE_TOOL_Brush3DPart, PARTICLE_TOOL_Brush3DNoise,
         PARTICLE_TOOL_Brush3DAttract, PARTICLE_TOOL_Brush3DWeight,
         PARTICLE_TOOL_Brush3DClone, PARTICLE_TOOL_Brush3DColor)

classes = [
    PARTICLE_OT_HairBrush3DEdit, PARTICLE_OT_HairBrush3DCloneRotate,
    VIEW3D_PT_tools_hair_brush_3d_options, PARTICLE_UL_color_group,
    VIEW3D_PT_hair_color_options, PARTICLE_OT_DeleteColorGroup,
    PARTICLE_OT_AssignColor2Selected, PARTICLE_OT_SelectInGroup,
    PARTICLE_OT_UndoDelColorGroup, HairBrush3DCommonProps, HairColorGroup,
    PARTICLE_OT_BrushCurvePreset, VIEW3D_PT_tools_particle_brush_falloff
]


def register():
    from bpy.utils import register_class, register_tool
    from . import particle_tools
    patch_VIEW3D_PT_tools_active()
    for cls in classes:
        register_class(cls)

    register_icons()
    for i, tool in enumerate(tools):
        BrushType.add(tool)
        register_tool(tool, separator=bool(i == 0))
    pie_menu.register()
    particle_tools.register()
    bpy.app.handlers.save_pre.append(hair_color_saving)
    bpy.app.handlers.load_post.append(hair_color_loading)
    bpy.app.handlers.depsgraph_update_post.append(sync_hair_color_on_deps_update)
    bpy.msgbus.subscribe_rna(
        key=(bpy.types.Object, "mode"),
        owner=particle_hook,
        args=(),
        notify=on_object_mode_change,
        options={'PERSISTENT'})

    particle_hook.set_particle_edit_tool_update_func(
        bpy.types.ParticleEdit.bl_rna.properties['tool'].as_pointer(), 1)
    particle_hook.set_mesh_calc_normals_function(
        bpy.types.Mesh.bl_rna.functions['calc_normals'].as_pointer())
    bpy.msgbus.subscribe_rna(
        key=(bpy.types.ParticleEdit, "tool"),
        owner=particle_hook,
        args=(),
        notify=on_particle_edit_tool_change,
        options={'PERSISTENT'})


def unregister():
    from bpy.utils import unregister_class, unregister_tool
    from . import pie_menu
    from . import particle_tools
    restore_VIEW3D_PT_tools_active()
    for cls in reversed(classes):
        unregister_class(cls)

    for tool in tools:
        unregister_tool(tool)
    pie_menu.unregister()
    particle_tools.unregister()
    bpy.app.handlers.save_pre.remove(hair_color_saving)
    bpy.app.handlers.load_post.remove(hair_color_loading)
    bpy.app.handlers.depsgraph_update_post.remove(sync_hair_color_on_deps_update)
    particle_hook.set_particle_edit_tool_update_func(
        bpy.types.ParticleEdit.bl_rna.properties['tool'].as_pointer(), 0)
    bpy.msgbus.clear_by_owner(particle_hook)


if __name__ == "__main__":
    register()
