# ##### 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 #####

from math import sin, cos, pi, fabs
import bpy
import bgl
import gpu
from mathutils import Vector, kdtree
from bpy_extras import view3d_utils
from bpy.app.handlers import persistent
from .particle_hook import particle_hook
from .draw_hair import show_hair_color, update_hair_color, ColorUpdateFlag

UNIT_VEC = Vector((1, 1, 1)).normalized()
is_293 = bpy.app.version >= (2, 93)


def gl_viewport(x, y, width, height):
    if is_293:
        gpu.state.viewport_set(x, y, width, height)
    else:
        bgl.glViewport(x, y, width, height)


def gl_line_width(line_width):
    if is_293:
        gpu.state.line_width_set(line_width)
    else:
        bgl.glLineWidth(line_width)


def gl_blend_alpha(enable):
    if is_293:
        gpu.state.blend_set('ALPHA' if enable else 'NONE')
    else:
        if enable:
            bgl.glEnable(bgl.GL_BLEND)
        else:
            bgl.glDisable(bgl.GL_BLEND)


def gl_line_smooth(enable):
    if not is_293:
        if enable:
            bgl.glEnable(bgl.GL_LINE_SMOOTH)
        else:
            bgl.glDisable(bgl.GL_LINE_SMOOTH)


def get_circle_wire_2d(x, y, segs):
    mul = (1.0 / (segs - 1)) * (pi * 2)
    return [(sin(i * mul) + x, cos(i * mul) + y) for i in range(segs)]


def location_3d_to_vector_3d(rv3d: bpy.types.RegionView3D, loc: Vector) -> Vector:
    if rv3d.is_perspective:
        p2 = rv3d.view_matrix @ Vector((*loc, 1.0))
        p2.xyz *= 2.0
        p2 = rv3d.view_matrix.inverted() @ p2
        return (loc - p2.xyz).normalized()
    return rv3d.view_matrix.inverted().col[2].xyz.normalized()


def calc_projected_radius(obmat, region, rv3d, rad, loc):
    view = location_3d_to_vector_3d(rv3d, loc)
    nonortho = view.copy()
    if fabs(nonortho.x) < 0.1:
        nonortho.x += 1.0
    elif fabs(nonortho.y) < 0.1:
        nonortho.y += 1.0
    else:
        nonortho.z += 1.0
    ortho = nonortho.cross(view).normalized()
    offset = loc + ortho * rad
    p1 = view3d_utils.location_3d_to_region_2d(region, rv3d, loc)
    p2 = view3d_utils.location_3d_to_region_2d(region, rv3d, offset)
    if p1 and p2:
        scale = (obmat.to_3x3() @ UNIT_VEC).length
        if scale == 0:
            scale = 1.0
        return int((p1 - p2).length * scale)
    return 0


def calc_unprojected_radius(obmat, loc, region, rv3d, pixel_rad):
    world_loc = obmat @ loc
    pers_mat = rv3d.perspective_matrix.copy()
    zfac = fabs(world_loc.dot(pers_mat[3].xyz) + pers_mat[3][3])
    if zfac < 1e-6:
        zfac = 1.0
    dx = 2.0 * pixel_rad * zfac / region.width
    pers_mat.invert()
    delta = Vector((pers_mat[0][0] * dx, pers_mat[1][0] * dx, pers_mat[2][0] * dx))
    radius3d = delta.length

    scale = (obmat.to_3x3() @ UNIT_VEC).length
    if scale == 0:
        scale = 1.0
    return radius3d, radius3d / scale


def isect_object_3d_ray(obj, ray_origin, ray_target):
    # cast rays and find the closest object
    obmat = obj.matrix_world.copy()

    # get the ray relative to the object
    matrix_inv = obmat.inverted()
    ray_origin_obj = matrix_inv @ ray_origin
    ray_target_obj = matrix_inv @ ray_target
    ray_direction_obj = ray_target_obj - ray_origin_obj

    # cast the ray
    on_mesh, loc, nor, _ = obj.ray_cast(ray_origin_obj, ray_direction_obj)
    return on_mesh, loc, nor


def isect_object_from_2d_coord(region, rv3d, obj, xy):
    # get the context arguments
    coord = xy[0] - region.x, xy[1] - region.y
    clamp = None
    if not rv3d.is_perspective and rv3d.view_perspective != 'CAMERA':
        clamp = rv3d.view_distance * 2

    # get the ray from the viewport and mouse
    view_vector = view3d_utils.region_2d_to_vector_3d(region, rv3d, coord)
    ray_origin = view3d_utils.region_2d_to_origin_3d(region, rv3d, coord, clamp=clamp)
    ray_target = ray_origin + view_vector

    return isect_object_3d_ray(obj, ray_origin, ray_target)


def get_addon_name():
    pn = __package__
    di = pn.find(".")
    if di != -1:
        pn = pn[:di]
    return pn


class HairRootsTree:
    hair_roots = None
    tree_size = -1
    psys = 0

    @staticmethod
    def generate_hair_roots(psys_pointer, obj_eval):
        HairRootsTree.psys = psys_pointer
        psys_eval = obj_eval.particle_systems.active
        HairRootsTree.tree_size = len(psys_eval.particles)
        HairRootsTree.hair_roots = kdtree.KDTree(HairRootsTree.tree_size)
        for i, p in enumerate(psys_eval.particles):
            HairRootsTree.hair_roots.insert(p.hair_keys[0].co, i)
        HairRootsTree.hair_roots.balance()

    @staticmethod
    def use_exists_roots_tree(psys_pointer, ktree, size):
        HairRootsTree.psys = psys_pointer
        ktree.balance()
        HairRootsTree.hair_roots = ktree
        HairRootsTree.tree_size = size

    @staticmethod
    def reset_roots_tree():
        HairRootsTree.hair_roots = None
        HairRootsTree.tree_size = -1
        HairRootsTree.psys = 0


class ColorGroupManager:
    undo_stack = []

    @staticmethod
    def generate_color_groups(psys: bpy.types.ParticleSystem):
        colors = particle_hook.get_hair_color_groups(psys.as_pointer())
        pset = psys.settings
        groups = pset.hair_color_groups
        groups.clear()
        j = 0
        for i in range(0, len(colors), 3):
            group = groups.add()
            group.rgb = (colors[i], colors[i + 1], colors[i + 2])
            group.name = str(j)
            j += 1

    @staticmethod
    def modify_color_groups(operator, context, rgb, op_code, psys, draw_flag):
        if len(psys.settings['hair_color__']) != len(psys.particles) * 3:
            operator.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

        particle_hook.set_hair_color_group(psys.as_pointer(), rgb, op_code)
        ColorGroupManager.generate_color_groups(psys)
        update_hair_color(context, draw_flag)

    @staticmethod
    def clear_color_groups(pset, draw_flag, context):
        if "hair_color__" in pset:
            del pset["hair_color__"]
            if draw_flag is not None:
                context = context if context else bpy.context
                update_hair_color(context, draw_flag)
        pset.hair_color_groups.clear()

    @staticmethod
    def on_add_hair(obj_eval, added_length, psys: bpy.types.ParticleSystem):
        ColorGroupManager.undo_stack = []
        particle_hook.extend_color_array(psys.as_pointer(), added_length)
        ColorGroupManager.generate_color_groups(psys)
        HairRootsTree.generate_hair_roots(psys.as_pointer(), obj_eval)

    @staticmethod
    def on_delete_hair(obj_eval, psys):
        ColorGroupManager.undo_stack = []
        psys_eval = obj_eval.particle_systems.active
        find_root_func = HairRootsTree.hair_roots.find
        count = len(psys_eval.particles)
        new_roots = kdtree.KDTree(count)
        if count == 0:
            if "hair_color__" in psys.settings:
                del psys.settings["hair_color__"]
        else:
            exist_index = []
            for i, p in enumerate(psys_eval.particles):
                new_roots.insert(p.hair_keys[0].co, i)
                co, index, dist = find_root_func(p.hair_keys[0].co)
                exist_index.append(index)
            particle_hook.shrink_color_array(psys.as_pointer(), exist_index)

        old_group_count = len(psys.settings.hair_color_groups)
        ColorGroupManager.generate_color_groups(psys)
        HairRootsTree.use_exists_roots_tree(psys.as_pointer(), new_roots, count)
        if len(psys.settings.hair_color_groups) != old_group_count:
            for area in bpy.context.screen.areas:
                if area.type == "VIEW_3D":
                    area.tag_redraw()
                    return

    @classmethod
    def undo_hair_color(cls, psys, context):
        backup = ColorGroupManager.undo_stack.pop()
        backup.pop()
        if len(psys.particles) * 3 == len(backup):
            psys.settings["hair_color__"] = backup
        else:
            r, g, b = backup.pop()
            colors = psys.settings["hair_color__"]
            for i in backup:
                colors[i] = r
                colors[i + 1] = g
                colors[i + 2] = b

        ColorGroupManager.generate_color_groups(psys)
        update_hair_color(context, ColorUpdateFlag.FORCE_TEX)

    @classmethod
    def backup_hair_color(cls, pset, rgb=(1.0, 1.0, 1.0), backup_all=False):
        if "hair_color__" not in pset:
            return

        colors = pset["hair_color__"]
        total = len(colors)
        if backup_all:
            backup = list(colors)
        else:
            backup = []
            append = backup.append
            for i in range(0, total, 3):
                if colors[i] == rgb[0] and colors[i + 1] == rgb[1] and colors[i + 2] == rgb[2]:
                    append(i)
            append(rgb)

        backup.append(total)
        ColorGroupManager.undo_stack.append(backup)


@persistent
def hair_color_saving(scene, deps):
    need_clear_color = False
    for sc in bpy.data.scenes:
        for obj in sc.objects:
            for psys in obj.particle_systems:
                pset = psys.settings
                if len(pset.hair_color_groups) == 1:
                    need_clear_color = True
                    ColorGroupManager.clear_color_groups(pset, None, None)
    if need_clear_color:
        update_hair_color(bpy.context, ColorUpdateFlag.CLEAR_TEX)


@persistent
def hair_color_loading(scene, deps):
    for sc in bpy.data.scenes:
        for obj in sc.objects:
            for psys in obj.particle_systems:
                pset = psys.settings
                hair_count = len(psys.particles)
                if "hair_color__" in pset:
                    if hair_count * 3 != len(pset["hair_color__"]):
                        ColorGroupManager.clear_color_groups(pset, None, None)
    show_hair_color(bpy.context, bpy.context.scene.hair_brush_3d.show_hair_color)
    update_hair_color(bpy.context, ColorUpdateFlag.LOADED)


@persistent
def sync_hair_color_on_deps_update(scene, deps):
    obj = scene.tool_settings.particle_edit.object
    update_hair_color(bpy.context, ColorUpdateFlag.FROM_DEPS)
    if not obj or not obj.particle_systems.active:
        return
    psys = obj.particle_systems.active
    pset = psys.settings
    if pset.type != 'HAIR':
        return
    if not psys.is_edited:
        ColorGroupManager.clear_color_groups(psys.settings, ColorUpdateFlag.CLEAR_TEX | ColorUpdateFlag.FROM_DEPS, None)
        if obj.mode == 'PARTICLE_EDIT':
            ColorGroupManager.undo_stack = []
            HairRootsTree.generate_hair_roots(psys.as_pointer(), obj.evaluated_get(deps))
        return
    if obj.mode != 'PARTICLE_EDIT':
        return
    if psys.as_pointer() != HairRootsTree.psys:
        HairRootsTree.generate_hair_roots(psys.as_pointer(), obj.evaluated_get(deps))
        ColorGroupManager.undo_stack = []
        return
    if "hair_color__" not in pset:
        if HairRootsTree.tree_size != len(psys.particles):
            HairRootsTree.generate_hair_roots(psys.as_pointer(), obj.evaluated_get(deps))
        return
    length_needed = len(psys.particles) * 3
    color_length = len(pset["hair_color__"])
    if length_needed < color_length:
        ColorGroupManager.on_delete_hair(obj.evaluated_get(deps), psys)
        update_hair_color(bpy.context, ColorUpdateFlag.FORCE_TEX | ColorUpdateFlag.FROM_DEPS)
    elif length_needed > color_length:
        ColorGroupManager.on_add_hair(obj.evaluated_get(deps), length_needed - color_length, psys)
