# ##### 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 contextlib
import bpy
import bpy.app.handlers as handlers
from .hair_modifier import _hair_modifier
from .utils import get_enum_items_item, tag_view3d_redraw
from .props import HairModifierProps


def get_evaluated_id(deps, id):
    if id.is_evaluated:
        return id
    return id.evaluated_get(deps)


def get_original_id(id):
    if id.is_evaluated:
        return id.original
    return id


class UpdateType:
    NONE = 0
    ADD = 1
    MOVE = 2
    ACTIVE_CHANGE = 3
    REMOVE = 4
    MODIFIED = 5


class ApplyFlag:
    ADD = 1
    MOVE = 2
    REMOVE = 4
    DO_RENDER = 8
    EXPORT = 16


class HairParticleSystem:
    def __init__(self, psys) -> None:
        self.edit_updated = False
        self.edit_updated_apply_args = 0
        self.modifiers = []
        self.sync_modifiers(psys)

    def sync_modifiers(self, psys):
        part = get_original_id(psys.settings)
        if len(self.modifiers) == len(part.child_modifier.modifiers):
            return False
        self.free_modifiers()
        _hair_modifier.reset_child_pathcache(psys.as_pointer())
        for hmod in part.child_modifier.modifiers:
            self.modifiers.append(_hair_modifier.create_modifier(
                get_enum_items_item(HairModifierProps.MODIFIER_TYPE_ITEMS, hmod.type, 4)))
        return True

    def add_modifier(self, type):
        self.modifiers.append(_hair_modifier.create_modifier(
            get_enum_items_item(HairModifierProps.MODIFIER_TYPE_ITEMS, type, 4)))

    def remove_modifier(self, index):
        assert 0 <= index < len(self.modifiers)
        mod = self.modifiers.pop(index)
        _hair_modifier.free_modifier(mod)

    def move_modifier(self, orig_idx, tar_idx):
        assert 0 <= orig_idx < len(self.modifiers)
        assert 0 <= tar_idx < len(self.modifiers)
        self.modifiers[orig_idx], self.modifiers[tar_idx] = self.modifiers[tar_idx], self.modifiers[orig_idx]

    def apply_modifiers(self, obj, psys, flag=0):
        self.sync_modifiers(psys)
        if obj.mode == 'PARTICLE_EDIT' and obj.particle_systems.active.name == psys.name:
            self.edit_updated = True
            self.edit_updated_apply_args = (obj, psys, flag)
            return
        _hair_modifier.apply_modifiers(self.modifiers, obj.as_pointer(), psys.as_pointer(), flag)

    def get_modifier_pointer(self, index):
        if index >= 0 and index < len(self.modifiers):
            return self.modifiers[index]
        return None

    def edit_apply_modifiers(self):
        if self.edit_updated:
            obj, psys, flag = self.edit_updated_apply_args
            _hair_modifier.apply_modifiers(self.modifiers, obj.as_pointer(), psys.as_pointer(), flag)
            self.edit_updated = False

    def free_modifiers(self):
        for mod in self.modifiers:
            _hair_modifier.free_modifier(mod)
        self.modifiers.clear()

    def __del__(self):
        self.free_modifiers()


class HairObject:
    def __init__(self) -> None:
        self.psyss = {}

    def update_particle_systems(self, obj, is_render):
        if obj.mode not in ('OBJECT', 'PARTICLE_EDIT'):
            return

        pre_keys = set(self.psyss.keys())
        cur_keys = set()
        apply_flag = 0
        if is_render:
            apply_flag = ApplyFlag.DO_RENDER
        for mod in obj.modifiers:
            if (mod.type != 'PARTICLE_SYSTEM' or (not is_render and not mod.show_viewport)
                    or (is_render and not mod.show_render)):
                continue
            psys: bpy.types.ParticleSystem = mod.particle_system
            part = get_original_id(psys.settings)
            if (part.type != 'HAIR' or len(psys.child_particles) == 0 or
                part.display_method != 'RENDER' or part.render_type != 'PATH'):
                continue
            if len(part.child_modifier.modifiers) == 0:
                if psys.name in self.psyss:
                    self.psyss[psys.name].sync_modifiers(psys)
                continue

            cur_keys.add(psys.name)
            if psys.name not in self.psyss:
                hpsys = HairParticleSystem(psys)
                hpsys.apply_modifiers(obj, psys, apply_flag)
                self.psyss[psys.name] = hpsys
            else:
                hpsys = self.psyss[psys.name]
                hpsys.sync_modifiers(psys)

        diff = pre_keys.difference(cur_keys)
        for name in diff:
            del self.psyss[name]


class HairModifierManager:
    hair_objs = {}
    render_hair_objs = {}
    obj_count = -1
    render_obj_count = -1
    is_exporting_abc = False
    is_block_update = False
    scene_draw_pre_handler = None

    @staticmethod
    def scene_pre_draw_cb():
        obj = bpy.context.object
        if obj and obj.mode == 'PARTICLE_EDIT' and obj in HairModifierManager.hair_objs:
            hair_obj = HairModifierManager.hair_objs[obj]
            if obj.particle_systems.active and obj.particle_systems.active.name in hair_obj.psyss:
                hair_obj.psyss[obj.particle_systems.active.name].edit_apply_modifiers()

    @classmethod
    def update_object(cls, orig_obj, eval_obj, is_render):
        if eval_obj.type != 'MESH' or len(eval_obj.particle_systems) == 0:
            return None
        if is_render:
            hair_objs = cls.render_hair_objs
        else:
            hair_objs = cls.hair_objs
        if orig_obj in hair_objs:
            hair_obj = hair_objs[orig_obj]
            hair_obj.update_particle_systems(eval_obj, is_render)
        else:
            hair_obj = HairObject()
            hair_obj.update_particle_systems(eval_obj, is_render)
            if hair_obj.psyss:
                hair_objs[orig_obj] = hair_obj
        return hair_obj

    @classmethod
    def get_hair_modifier_pointer(cls, obj, psys, mod_idx):
        if obj in cls.hair_objs:
            hair_obj = cls.hair_objs[obj]
            if psys.name in hair_obj.psyss:
                hpsys = hair_obj.psyss[psys.name]
                return hpsys.get_modifier_pointer(mod_idx)
        return None

    @classmethod
    def update_objects(cls, scene, deps, force=False):
        is_render = deps.mode == 'RENDER'
        ups = len(deps.objects)
        if is_render:
            if cls.render_obj_count == ups and not force:
                return
            cls.render_obj_count = ups
            hair_objs = cls.render_hair_objs
        else:
            if cls.render_hair_objs:
                cls.render_hair_objs.clear()
            cls.render_obj_count = -1
            if bpy.context.mode == 'PARTICLE':
                if cls.scene_draw_pre_handler is None:
                    cls.scene_draw_pre_handler = bpy.types.SpaceView3D.draw_handler_add(
                        HairModifierManager.scene_pre_draw_cb, (), 'WINDOW', 'PRE_VIEW')
            elif cls.scene_draw_pre_handler:
                bpy.types.SpaceView3D.draw_handler_remove(cls.scene_draw_pre_handler, 'WINDOW')
                cls.scene_draw_pre_handler = None
            if cls.obj_count == ups and not force:
                return
            cls.obj_count = ups
            hair_objs = cls.hair_objs

        pre_keys = set(hair_objs.keys())
        cur_keys = set()
        for obj in deps.objects:
            orig_obj = get_original_id(obj)
            hair_obj = cls.update_object(orig_obj, obj, is_render)
            if hair_obj and hair_obj.psyss:
                cur_keys.add(orig_obj)

        diff = pre_keys.difference(cur_keys)
        for key in diff:
            del hair_objs[key]

    @classmethod
    def apply_updated(cls, obj, psys, update_type, extra_info):
        if HairModifierManager.is_block_update or not psys or psys.settings.type != 'HAIR':
            return
        orig_obj = get_original_id(obj)
        if orig_obj not in cls.hair_objs:
            hair_obj = HairObject()
            hair_obj.update_particle_systems(obj, False)
            cls.hair_objs[orig_obj] = hair_obj
        else:
            hair_obj = cls.hair_objs[orig_obj]
            flag = 0
            if psys.name not in hair_obj.psyss:
                hair_obj.update_particle_systems(obj, False)
                flag = ApplyFlag.ADD
            else:
                hpsys = hair_obj.psyss[psys.name]
                if update_type == UpdateType.ADD:
                    hpsys.add_modifier(extra_info)
                    flag = ApplyFlag.ADD
                elif update_type == UpdateType.REMOVE:
                    hpsys.remove_modifier(extra_info)
                    flag = ApplyFlag.REMOVE
                elif update_type == UpdateType.MOVE:
                    hpsys.move_modifier(*extra_info)
                    flag = ApplyFlag.MOVE
                hpsys.apply_modifiers(obj, psys, flag)
            if flag:
                cls.render_hair_objs.clear()
                cls.render_obj_count = -1

    @classmethod
    def apply_all_objects(cls, deps, update=False):
        if deps.mode == 'RENDER':
            apply_flag = ApplyFlag.DO_RENDER
            hair_objs = cls.render_hair_objs
        else:
            apply_flag = 0
            hair_objs = cls.hair_objs
        if HairModifierManager.is_exporting_abc:
            apply_flag |= ApplyFlag.EXPORT
        for obj, hair_obj in hair_objs.items():
            eval_obj = get_evaluated_id(deps, obj)
            if update:
                hair_obj.update_particle_systems(eval_obj, deps.mode == 'RENDER')
            for psys in eval_obj.particle_systems:
                if psys.name in hair_obj.psyss:
                    hair_obj.psyss[psys.name].apply_modifiers(eval_obj, psys, apply_flag)

    @classmethod
    def clear_objects(cls, deps):
        if deps:
            if deps.mode == 'RENDER':
                cls.render_hair_objs.clear()
                cls.render_obj_count = -1
            else:
                cls.hair_objs.clear()
                cls.obj_count = -1
        else:
            cls.hair_objs.clear()
            cls.render_hair_objs.clear()
            cls.obj_count = -1
            cls.render_obj_count = -1


@contextlib.contextmanager
def block_update():
    HairModifierManager.is_block_update = True
    yield None
    HairModifierManager.is_block_update = False


def get_clump_point_count(obj, psys, mod_idx):
    if obj and psys:
        hmd_p = HairModifierManager.get_hair_modifier_pointer(obj, psys, mod_idx)
        if hmd_p:
            return int(_hair_modifier.get_clump_modifier_clump_point_count(hmd_p))
    return 0


def modifier_updated(context, settings: bpy.types.ParticleSettings, update_type, extra_info):
    if update_type == UpdateType.ADD:
        settings.child_modifier.modifiers[-1].touch_props()
    if settings.users == 1:
        if context.object is None or context.object.mode not in ('OBJECT', 'PARTICLE_EDIT'):
            return
        obj = get_evaluated_id(context.evaluated_depsgraph_get(), context.object)
        psys = obj.particle_systems.active
        if psys is None:
            return
        HairModifierManager.apply_updated(obj, psys, update_type, extra_info)
    elif settings.users > 1:
        for obj in context.visible_objects:
            if obj.type != 'MESH' or obj.mode not in ('OBJECT', 'PARTICLE_EDIT'):
                continue
            obj = get_evaluated_id(context.evaluated_depsgraph_get(), obj)
            for psys in obj.particle_systems:
                if get_original_id(psys.settings) == settings:
                    HairModifierManager.apply_updated(obj, psys, update_type, extra_info)


@handlers.persistent
def depsgraph_updated(scene, deps):
    HairModifierManager.update_objects(scene, deps)
    is_render = deps.mode == 'RENDER'
    ups = len(deps.updates)
    if ups == 0:
        obj = getattr(bpy.context, "object", None)
        if obj is None:
            return
        obj = obj.evaluated_get(deps)
        HairModifierManager.apply_updated(obj, obj.particle_systems.active, UpdateType.MODIFIED, None)
        tag_view3d_redraw(bpy.context, False)
    elif deps.id_type_updated('OBJECT'):
        for update in deps.updates:
            if isinstance(update.id, bpy.types.Object):
                obj = update.id
                hair_obj = HairModifierManager.update_object(get_original_id(obj), obj, is_render)
                for psys in obj.particle_systems:
                    if psys.name in hair_obj.psyss:
                        hair_obj.psyss[psys.name].apply_modifiers(obj, psys)
    elif deps.id_type_updated('PARTICLE'):
        obj = getattr(bpy.context, "object", None)
        if obj and ups == 1:
            settings = obj.particle_systems.active.settings
            if settings.users == 1:
                eval_obj = get_evaluated_id(deps, obj)
                HairModifierManager.apply_updated(eval_obj, eval_obj.particle_systems.active, UpdateType.MODIFIED, None)
            elif settings.users > 1:
                for object_instance in deps.object_instances:
                    obj = object_instance.object
                    if object_instance.is_instance or obj.type != 'MESH' or obj.mode not in ('OBJECT', 'PARTICLE_EDIT'):
                        continue
                    obj = get_evaluated_id(deps, obj)
                    for psys in obj.particle_systems:
                        if get_original_id(psys.settings) == settings:
                            HairModifierManager.apply_updated(obj, psys, UpdateType.MODIFIED, None)
        else:
            for update in deps.updates:
                if isinstance(update.id, bpy.types.ParticleSettings):
                    modifier_updated(bpy.context, update.id, UpdateType.MODIFIED, None)


@handlers.persistent
def frame_change_posted(scene, deps):
    HairModifierManager.update_objects(scene, deps)
    try:
        HairModifierManager.apply_all_objects(deps)
    except ReferenceError:
        HairModifierManager.clear_objects(deps)
    if HairModifierManager.is_exporting_abc and bpy.context.window_manager.is_interface_locked is False:
        HairModifierManager.is_exporting_abc = False
        _hair_modifier.restore_scene_cddata_masks(get_original_id(scene).as_pointer())


@handlers.persistent
def undo_posted(scene):
    deps = bpy.context.evaluated_depsgraph_get()
    HairModifierManager.update_objects(scene, deps, force=True)
    HairModifierManager.apply_all_objects(deps, update=True)


@handlers.persistent
def load_pre(scene, deps):
    HairModifierManager.clear_objects(None)


@handlers.persistent
def load_post(scene, deps):
    undo_posted(scene)


def register_handlers():
    handlers.depsgraph_update_post.append(depsgraph_updated)
    handlers.frame_change_post.append(frame_change_posted)
    handlers.load_pre.append(load_pre)
    handlers.undo_post.append(undo_posted)
    handlers.load_post.append(load_post)


def unregister_handlers():
    handlers.depsgraph_update_post.remove(depsgraph_updated)
    handlers.frame_change_post.remove(frame_change_posted)
    handlers.load_pre.remove(load_pre)
    handlers.undo_post.remove(undo_posted)
    handlers.load_post.remove(load_post)
