# SPDX-FileCopyrightText: 2015 Blender Studio Tools Authors
#
# SPDX-License-Identifier: GPL-3.0-or-later

# Inspired by https://animplay.wordpress.com/2015/11/18/smear-frame-script-maya/.

# This addon allows the user to specify a camera and a collection,
# and create a 2D lattice that fills the camera's view,
# to deform the mesh objects in that collection.

import bpy
import math
from bpy.app.handlers import persistent
from mathutils import Vector
from bpy.props import (
    BoolProperty,
    PointerProperty,
    CollectionProperty,
    IntProperty,
    EnumProperty,
    FloatProperty,
)
from mathutils.geometry import intersect_point_line

from .utils import bounding_box_center_of_objects
from .prefs import get_addon_prefs


class CAMLAT_UL_lattice_slots(bpy.types.UIList):
    def draw_item(self, context, layout, data, item, icon, active_data, active_propname):
        lattice_slots = context.scene.lattice_slots
        active_slot = lattice_slots[context.scene.active_lattice_index]
        current_slot = item

        if self.layout_type in {'DEFAULT', 'COMPACT'}:
            if current_slot.collection:
                row = layout.row()
                icon = 'OUTLINER_COLLECTION' if current_slot.enabled else 'COLLECTION_COLOR_07'
                row.prop(current_slot.collection, 'name', text="", emboss=False, icon=icon)
                row.enabled = current_slot.enabled

                layout.prop(current_slot, 'strength', text="", slider=True, emboss=False)

                icon = 'CHECKBOX_HLT' if current_slot.enabled else 'CHECKBOX_DEHLT'
                layout.prop(current_slot, 'enabled', text="", icon=icon, emboss=False)
            else:
                layout.label(text="", translate=False, icon='COLLECTION_NEW')
        elif self.layout_type in {'GRID'}:
            layout.alignment = 'CENTER'
            layout.label(text="", icon_value=icon)


class LatticeSlot(bpy.types.PropertyGroup):
    enabled: BoolProperty(
        name="Enabled", description="Whether the Lattice has an effect or not", default=True
    )
    strength: FloatProperty(
        name="Strength", description="Strength of the lattice effect", min=0, max=1, default=1
    )
    lattice: PointerProperty(
        name="Lattice",
        type=bpy.types.Object,
        description="Lattice object generated by this LatticeSlot. This cannot be specified manually, use the Generate or Delete operator below",
    )

    def is_camera(self, obj):
        return obj.type == 'CAMERA'

    camera: PointerProperty(
        name="Camera",
        type=bpy.types.Object,
        description="Camera used by this LatticeSlot",
        poll=is_camera,
    )
    collection: PointerProperty(
        name="Collection",
        type=bpy.types.Collection,
        description="Collection affected by this LatticeSlot",
    )

    resolution: IntProperty(
        name="Resolution",
        description="Resolution of the lattice grid",
        min=5,
        max=64,
        default=10,
        options=set(),
    )


class OBJECT_OT_camlattice_add(bpy.types.Operator):
    """Add a Camera Lattice Slot"""

    bl_idname = "lattice.add_slot"
    bl_label = "Add Lattice Slot"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    def execute(self, context):
        scene = context.scene
        lattice_slots = scene.lattice_slots
        active_index = scene.active_lattice_index
        to_index = active_index + 1
        if len(lattice_slots) == 0:
            to_index = 0

        scene.lattice_slots.add()
        scene.lattice_slots.move(len(scene.lattice_slots) - 1, to_index)
        scene.active_lattice_index = to_index

        return {'FINISHED'}


class OBJECT_OT_camlattice_remove(bpy.types.Operator):
    """Remove Lattice Slot along with its Lattice object, animation and modifiers"""

    bl_idname = "lattice.remove_slot"
    bl_label = "Remove Lattice Slot"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    index: IntProperty()

    @classmethod
    def poll(cls, context):
        scene = context.scene
        if len(scene.lattice_slots) > 0:
            return True
        
        cls.poll_message_set("No slots to remove.")
        return False

    def execute(self, context):
        scene = context.scene
        lattice_slots = scene.lattice_slots
        active_index = scene.active_lattice_index
        # This behaviour is inconsistent with other UILists in Blender, but I am right and they are wrong!
        active_slot = lattice_slots[active_index]
        if active_slot.lattice:
            bpy.ops.lattice.delete_lattice_from_slot()
        to_index = active_index
        if to_index > len(lattice_slots) - 2:
            to_index = len(lattice_slots) - 2

        scene.lattice_slots.remove(self.index)
        scene.active_lattice_index = to_index

        return {'FINISHED'}


class OBJECT_OT_camlattice_move(bpy.types.Operator):
    """Move Lattice Slot"""

    bl_idname = "lattice.move_slot"
    bl_label = "Move Lattice Slot"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    direction: EnumProperty(
        name="Direction",
        items=[
            ('UP', 'UP', 'UP'),
            ('DOWN', 'DOWN', 'DOWN'),
        ],
        default='UP',
    )

    @classmethod
    def poll(cls, context):
        scene = context.scene
        if len(scene.lattice_slots) > 1:
            return True
        
        cls.poll_message_set("No slots to re-order.")
        return False

    def execute(self, context):
        scene = context.scene
        lattice_slots = scene.lattice_slots
        active_index = scene.active_lattice_index
        to_index = active_index + (1 if self.direction == 'DOWN' else -1)

        if to_index > len(lattice_slots) - 1:
            to_index = 0
        if to_index < 0:
            to_index = len(lattice_slots) - 1

        scene.lattice_slots.move(active_index, to_index)
        scene.active_lattice_index = to_index

        return {'FINISHED'}


class OBJECT_OT_camlattice_generate(bpy.types.Operator):
    """Generate a lattice to smear the selected collection from the selected camera"""

    bl_idname = "lattice.generate_lattice_for_slot"
    bl_label = "Generate Lattice"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    @classmethod
    def poll(cls, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]

        if not active_slot.collection:
            cls.poll_message_set("A collection must be selected above.")
            return False
        if not active_slot.camera:
            cls.poll_message_set("A camera must be selected above.")
            return False
        if active_slot.lattice:
            cls.poll_message_set("This slot already has a lattice generated for it. You can delete it above if you wish to re-generate it.")
            return False
        return True


    def execute(self, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]

        collection = active_slot.collection
        camera = active_slot.camera
        resolution = active_slot.resolution

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

        # Create a lattice object.
        lattice_name = "Lattice_" + collection.name
        lattice = bpy.data.lattices.new(lattice_name)
        lattice_ob = bpy.data.objects.new(lattice_name, lattice)
        scene.collection.objects.link(lattice_ob)
        active_slot.lattice = lattice_ob
        bpy.ops.object.select_all(action='DESELECT')
        context.view_layer.objects.active = lattice_ob
        lattice_ob.select_set(True)

        # Align to camera (not really needed).
        lattice_ob.rotation_euler = camera.matrix_world.to_euler()

        # Parent to camera.
        lattice_ob.parent = camera
        lattice_ob.matrix_parent_inverse = camera.matrix_world.inverted()
        # Constrain to camera.
        constraint = lattice_ob.constraints.new('DAMPED_TRACK')
        constraint.target = camera
        constraint.track_axis = 'TRACK_Z'

        ### Placing the Lattice in the center of the camera's view, at the bounding box center of the collection's objects.
        # Find the bounding box center of the collection of objects
        all_meshes = [o for o in collection.all_objects if o.type == 'MESH']
        center = bounding_box_center_of_objects(all_meshes)

        # Define a line from the camera towards the camera's view direction
        cam_vec = Vector((0, 0, -1))  # Default aim vector of a camera (they point straight down)
        # Rotate the default vector by the camera's rotation
        cam_vec.rotate(camera.matrix_world.to_euler())
        cam_world_pos = camera.matrix_world.to_translation()
        cam_target_pos = cam_world_pos + cam_vec

        # Find the nearest point on this line to the bounding box center
        intersect = intersect_point_line(center, cam_world_pos, cam_target_pos)[0]

        # This is where the Lattice is placed!
        lattice_ob.location = intersect

        # Scale the lattice so that it fills up the camera's view
        # based on the distance of this point from the camera and the scene's aspect ratio.
        # https://fullpipeumbrella.com/en/blender-python-script-how-to-position/
        distance = (intersect - cam_world_pos).length
        fov = camera.data.angle

        scale_x = distance * math.sin(fov / 2) / math.cos(fov / 2) * 2
        aspect_ratio = (scene.render.resolution_x * scene.render.pixel_aspect_x) / (
            scene.render.resolution_y * scene.render.pixel_aspect_y
        )
        scale_y = scale_x / aspect_ratio
        lattice_ob.scale = [scale_x, scale_y, 1]

        # Set lattice resolution
        lattice.points_u = resolution
        lattice.points_v = round(resolution / aspect_ratio)
        lattice.points_w = 1

        # Create two shape keys.
        bpy.ops.lattice.smear_add_shape()
        bpy.ops.lattice.smear_add_shape()

        # Add Lattice modifiers
        for ob in all_meshes:
            # Skip those meshes which are already being deformed by another mesh in the same collection.
            skip = False
            for m in ob.modifiers:
                if m.type == 'MESH_DEFORM' and m.object in all_meshes:
                    skip = True
                    break
                if m.type == 'SURFACE_DEFORM' and m.target in all_meshes:
                    skip = True
                    break
            if not ob.visible_get():
                skip = True
            if skip:
                continue

            mod = ob.modifiers.new(name=lattice_ob.name, type='LATTICE')
            mod.object = lattice_ob
            # Add drivers for easy disabling
            index = len(scene.lattice_slots) - 1

            driver = ob.driver_add(f'modifiers["{lattice_ob.name}"].strength').driver
            driver.type = 'SUM'
            var = driver.variables.new()
            var.targets[0].id_type = 'SCENE'
            var.targets[0].id = scene
            var.targets[0].data_path = f'lattice_slots[{index}].strength'

            driver = ob.driver_add(f'modifiers["{lattice_ob.name}"].show_viewport').driver
            driver.type = 'SUM'
            var = driver.variables.new()
            var.targets[0].id_type = 'SCENE'
            var.targets[0].id = scene
            var.targets[0].data_path = f'lattice_slots[{index}].enabled'

            driver = ob.driver_add(f'modifiers["{lattice_ob.name}"].show_render').driver
            driver.type = 'SUM'
            var = driver.variables.new()
            var.targets[0].id_type = 'SCENE'
            var.targets[0].id = scene
            var.targets[0].data_path = f'lattice_slots[{index}].enabled'

        return {'FINISHED'}


class OBJECT_OT_camlattice_delete(bpy.types.Operator):
    """Delete Lattice object, its animation and modifiers that target it in the selected collection's objects"""

    bl_idname = "lattice.delete_lattice_from_slot"
    bl_label = "Delete Lattice"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    @classmethod
    def poll(cls, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]

        if not active_slot.lattice:
            cls.poll_message_set("This slot has no lattice to delete.")
            return False

        return True

    def execute(self, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]
        lattice_ob = active_slot.lattice
        lattice = lattice_ob.data

        # Delete modifiers and their drivers
        collection = active_slot.collection
        for ob in collection.all_objects:
            if not ob.type == 'MESH':
                continue
            for m in ob.modifiers[:]:
                if not (m.type == 'LATTICE' and m.object == lattice_ob):
                    continue
                ob.driver_remove(f'modifiers["{m.name}"].strength')
                ob.driver_remove(f'modifiers["{m.name}"].show_viewport')
                ob.driver_remove(f'modifiers["{m.name}"].show_render')
                ob.modifiers.remove(m)

        # Delete animation datablocks
        datablocks = [lattice, lattice_ob, lattice.shape_keys]
        for datablock in datablocks:
            if not datablock:
                continue
            if not datablock.animation_data:
                continue
            if not datablock.animation_data.action:
                continue
            bpy.data.actions.remove(datablock.animation_data.action)

        # Delte Lattice datablock
        bpy.data.objects.remove(lattice_ob)

        # Delete Object datablock
        bpy.data.lattices.remove(lattice)

        return {'FINISHED'}


class OBJECT_OT_camlattice_shapekey_add(bpy.types.Operator):
    """Add a shape key to the active Lattice Slot's lattice, named after the current frame number"""

    bl_idname = "lattice.smear_add_shape"
    bl_label = "Add Smear Shape"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    def execute(self, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]
        lattice_ob = active_slot.lattice
        lattice = lattice_ob.data

        name = "Basis"
        if lattice.shape_keys:
            name = "Frame " + str(scene.frame_current)
        lattice_ob.shape_key_add(name=name, from_mix=False)
        lattice_ob.active_shape_key_index = len(lattice.shape_keys.key_blocks) - 1
        block = lattice.shape_keys.key_blocks[-1]
        block.value = 1

        return {'FINISHED'}


def shape_key_poll(context):
    ob = context.object
    if not ob or ob.type != 'LATTICE':
        return False
    if not ob.data.shape_keys or len(ob.data.shape_keys.key_blocks) < 2:
        return False
    return True


class OBJECT_OT_camlattice_zero_all(bpy.types.Operator):
    """Set all shape key values to 0"""

    bl_idname = "lattice.shape_keys_zero_all"
    bl_label = "Zero All Shape Keys"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    @classmethod
    def poll(cls, context):
        return shape_key_poll(context)

    def execute(self, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]
        lattice_ob = active_slot.lattice
        lattice = lattice_ob.data

        for sk in lattice.shape_keys.key_blocks:
            sk.value = 0

        return {'FINISHED'}


class OBJECT_OT_camlattice_keyframe_all(bpy.types.Operator):
    """Insert a keyframe on the current frame for all shape key values"""

    bl_idname = "lattice.shape_keys_keyframe_all"
    bl_label = "Keyframe All Shape Keys"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    @classmethod
    def poll(cls, context):
        return shape_key_poll(context)

    def execute(self, context):
        scene = context.scene
        active_slot = scene.lattice_slots[scene.active_lattice_index]
        lattice_ob = active_slot.lattice
        lattice = lattice_ob.data

        for sk in lattice.shape_keys.key_blocks:
            sk.keyframe_insert('value')

        return {'FINISHED'}


class VIEW3D_PT_camlattice_panel(bpy.types.Panel):
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'Lattice Magic'
    bl_label = "Camera Lattice"

    @classmethod
    def poll(cls, context):
        return True

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False
        scene = context.scene
        active_index = scene.active_lattice_index

        row = layout.row()
        row.template_list(
            'CAMLAT_UL_lattice_slots',
            '',
            scene,
            'lattice_slots',
            scene,
            'active_lattice_index',
        )

        col = row.column()
        col.operator(OBJECT_OT_camlattice_add.bl_idname, text="", icon='ADD')
        remove_op = col.operator(OBJECT_OT_camlattice_remove.bl_idname, text="", icon='REMOVE')
        remove_op.index = active_index
        col.separator()
        move_up_op = col.operator(OBJECT_OT_camlattice_move.bl_idname, text="", icon='TRIA_UP')
        move_up_op.direction = 'UP'
        move_down_op = col.operator(OBJECT_OT_camlattice_move.bl_idname, text="", icon='TRIA_DOWN')
        move_down_op.direction = 'DOWN'

        if len(scene.lattice_slots) == 0:
            return

        active_slot = scene.lattice_slots[scene.active_lattice_index]

        col = layout.column()
        if active_slot.lattice:
            col.enabled = False
        row = col.row()
        if not active_slot.collection:
            row.alert = True
        row.prop(active_slot, 'collection')
        row = col.row()
        if not active_slot.camera:
            row.alert = True
        row.prop(active_slot, 'camera', icon='OUTLINER_OB_CAMERA')
        col.prop(active_slot, 'resolution')

        layout.separator()

        if not active_slot.lattice:
            layout.operator(OBJECT_OT_camlattice_generate.bl_idname, icon='OUTLINER_OB_LATTICE')
        else:
            layout.operator(OBJECT_OT_camlattice_delete.bl_idname, icon='TRASH')

        row = layout.row()
        row.enabled = False
        row.prop(active_slot, 'lattice')

        if not active_slot.lattice:
            return

        lattice_ob = active_slot.lattice
        lattice = lattice_ob.data

        col = layout.column()

        row = layout.row(align=True)
        row.operator(OBJECT_OT_camlattice_zero_all.bl_idname, text="", icon='RADIOBUT_OFF')
        row.operator(OBJECT_OT_camlattice_keyframe_all.bl_idname, text="", icon='HANDLETYPE_FREE_VEC')
        row.separator()
        prefs = get_addon_prefs(context)
        row.prop(prefs, 'update_active_shape_key', toggle=True, text="", icon='TIME')

        row = layout.row()

        # Display the lattice's Shape Keys in a less cluttered way than in the Properties editor.
        row.template_list(
            'MESH_UL_shape_keys',
            '',
            lattice.shape_keys,
            'key_blocks',
            lattice_ob,
            'active_shape_key_index',
        )

        col = row.column()
        col.operator(OBJECT_OT_camlattice_shapekey_add.bl_idname, text="", icon='ADD')
        remove_op = col.operator('object.shape_key_remove', text="", icon='REMOVE')

        col.separator()
        col.menu("MESH_MT_shape_key_context_menu", icon='DOWNARROW_HLT', text="")

        col.separator()
        move_up_op = col.operator('object.shape_key_move', text="", icon='TRIA_UP')
        move_up_op.type = 'UP'
        move_down_op = col.operator('object.shape_key_move', text="", icon='TRIA_DOWN')
        move_down_op.type = 'DOWN'


@persistent
def camera_lattice_frame_change(scene):
    """On frame change, set the active shape key of the active lattice object to the most recent frame

    (Assuming the shape keys are named after the frame on which they are used)
    """

    # I wonder why this function doesn't recieve a context... should it not be relied on from here? o.0
    context = bpy.context

    prefs = get_addon_prefs(context)
    if not prefs.update_active_shape_key:
        return

    ob = context.object
    if not shape_key_poll(context):
        return

    key_blocks = ob.data.shape_keys.key_blocks

    current_frame = scene.frame_current
    most_recent_number = 0
    most_recent_index = 1
    for i, kb in enumerate(key_blocks):
        if not kb.name.startswith('Frame '):
            continue
        number_str = kb.name[5:].split(".")[0]
        if number_str == "":
            continue
        number = int(number_str)
        if number <= current_frame and number >= most_recent_number:
            most_recent_number = number
            most_recent_index = i
            if most_recent_number == current_frame:
                break

    if ob.active_shape_key_index != most_recent_index:
        ob.active_shape_key_index = most_recent_index


registry = [
    LatticeSlot,
    CAMLAT_UL_lattice_slots,
    OBJECT_OT_camlattice_add,
    OBJECT_OT_camlattice_remove,
    OBJECT_OT_camlattice_move,
    OBJECT_OT_camlattice_generate,
    OBJECT_OT_camlattice_delete,
    OBJECT_OT_camlattice_shapekey_add,
    OBJECT_OT_camlattice_zero_all,
    OBJECT_OT_camlattice_keyframe_all,
    VIEW3D_PT_camlattice_panel,
]


def register():
    bpy.types.Scene.lattice_slots = CollectionProperty(type=LatticeSlot)
    bpy.types.Scene.active_lattice_index = IntProperty()
    bpy.app.handlers.frame_change_post.append(camera_lattice_frame_change)


def unregister():
    del bpy.types.Scene.lattice_slots
    del bpy.types.Scene.active_lattice_index

    bpy.app.handlers.frame_change_post.remove(camera_lattice_frame_change)