# SPDX-License-Identifier: GPL-2.0-or-later

import bpy
import gpu
import bgl
import blf
import math
import enum
import random

from itertools import islice
from mathutils.bvhtree import BVHTree
from mathutils import Vector, Matrix, Euler
from gpu_extras.batch import batch_for_shader

from bpy_extras.view3d_utils import (
    region_2d_to_vector_3d,
    region_2d_to_origin_3d
)


# Modal Operator
################################################################

class ScatterObjects(bpy.types.Operator):
    bl_idname = "object.scatter"
    bl_label = "Scatter Objects"
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        return (
            currently_in_3d_view(context)
            and context.active_object is not None
            and context.active_object.mode == 'OBJECT')

    def invoke(self, context, event):
        self.target_object = context.active_object
        self.objects_to_scatter = get_selected_non_active_objects(context)

        if self.target_object is None or len(self.objects_to_scatter) == 0:
            self.report({'ERROR'}, "Select objects to scatter and a target object")
            return {'CANCELLED'}

        self.base_scale = get_max_object_side_length(self.objects_to_scatter)

        self.targets = []
        self.active_target = None
        self.target_cache = {}

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

    def modal(self, context, event):
        context.area.tag_redraw()

        if not event_is_in_region(event, context.region) and self.active_target is None:
            return {'PASS_THROUGH'}

        if event.type == 'ESC':
            return self.finish('CANCELLED')

        if event.type == 'RET' and event.value == 'PRESS':
            self.create_scatter_object()
            return self.finish('FINISHED')

        event_used = self.handle_non_exit_event(event)
        if event_used:
            return {'RUNNING_MODAL'}
        else:
            return {'PASS_THROUGH'}

    def handle_non_exit_event(self, event):
        if self.active_target is None:
            if event.type == 'LEFTMOUSE' and event.value == 'PRESS':
                self.active_target = StrokeTarget()
                self.active_target.start_build(self.target_object)
                return True
        else:
            build_state = self.active_target.continue_build(event)
            if build_state == BuildState.FINISHED:
                self.targets.append(self.active_target)
                self.active_target = None
            self.remove_target_from_cache(self.active_target)
            return True

        return False

    def enable_draw_callback(self):
        self._draw_callback_view = bpy.types.SpaceView3D.draw_handler_add(self.draw_view, (), 'WINDOW', 'POST_VIEW')
        self._draw_callback_px = bpy.types.SpaceView3D.draw_handler_add(self.draw_px, (), 'WINDOW', 'POST_PIXEL')

    def disable_draw_callback(self):
        bpy.types.SpaceView3D.draw_handler_remove(self._draw_callback_view, 'WINDOW')
        bpy.types.SpaceView3D.draw_handler_remove(self._draw_callback_px, 'WINDOW')

    def draw_view(self):
        for target in self.iter_targets():
            target.draw()

        draw_matrices_batches(list(self.iter_matrix_batches()))

    def draw_px(self):
        draw_text((20, 20, 0), "Instances: " + str(len(self.get_all_matrices())))

    def finish(self, return_value):
        self.disable_draw_callback()
        bpy.context.area.tag_redraw()
        return {return_value}

    def create_scatter_object(self):
        matrix_chunks = make_random_chunks(
            self.get_all_matrices(), len(self.objects_to_scatter))

        collection = bpy.data.collections.new("Scatter")
        bpy.context.collection.children.link(collection)

        for obj, matrices in zip(self.objects_to_scatter, matrix_chunks):
            make_duplicator(collection, obj, matrices)

    def get_all_matrices(self):
        settings = self.get_current_settings()

        matrices = []
        for target in self.iter_targets():
            self.ensure_target_is_in_cache(target)
            matrices.extend(self.target_cache[target].get_matrices(settings))
        return matrices

    def iter_matrix_batches(self):
        settings = self.get_current_settings()
        for target in self.iter_targets():
            self.ensure_target_is_in_cache(target)
            yield self.target_cache[target].get_batch(settings)

    def iter_targets(self):
        yield from self.targets
        if self.active_target is not None:
            yield self.active_target

    def ensure_target_is_in_cache(self, target):
        if target not in self.target_cache:
            entry = TargetCacheEntry(target, self.base_scale)
            self.target_cache[target] = entry

    def remove_target_from_cache(self, target):
        self.target_cache.pop(self.active_target, None)

    def get_current_settings(self):
        return bpy.context.scene.scatter_properties.to_settings()

class TargetCacheEntry:
    def __init__(self, target, base_scale):
        self.target = target
        self.last_used_settings = None
        self.base_scale = base_scale
        self.settings_changed()

    def get_matrices(self, settings):
        self._handle_new_settings(settings)
        if self.matrices is None:
            self.matrices = self.target.get_matrices(settings)
        return self.matrices

    def get_batch(self, settings):
        self._handle_new_settings(settings)
        if self.gpu_batch is None:
            self.gpu_batch = create_batch_for_matrices(self.get_matrices(settings), self.base_scale)
        return self.gpu_batch

    def _handle_new_settings(self, settings):
        if settings != self.last_used_settings:
            self.settings_changed()
        self.last_used_settings = settings

    def settings_changed(self):
        self.matrices = None
        self.gpu_batch = None


# Duplicator Creation
######################################################

def make_duplicator(target_collection, source_object, matrices):
    triangle_scale = 0.1

    duplicator = triangle_object_from_matrices(source_object.name + " Duplicator", matrices, triangle_scale)
    duplicator.instance_type = 'FACES'
    duplicator.use_instance_faces_scale = True
    duplicator.show_instancer_for_viewport = True
    duplicator.show_instancer_for_render = False
    duplicator.instance_faces_scale = 1 / triangle_scale

    copy_obj = source_object.copy()
    copy_obj.name = source_object.name + " - copy"
    copy_obj.location = (0, 0, 0)
    copy_obj.parent = duplicator

    target_collection.objects.link(duplicator)
    target_collection.objects.link(copy_obj)

def triangle_object_from_matrices(name, matrices, triangle_scale):
    mesh = triangle_mesh_from_matrices(name, matrices, triangle_scale)
    return bpy.data.objects.new(name, mesh)

def triangle_mesh_from_matrices(name, matrices, triangle_scale):
    mesh = bpy.data.meshes.new(name)
    vertices, polygons = mesh_data_from_matrices(matrices, triangle_scale)
    mesh.from_pydata(vertices, [], polygons)
    mesh.update()
    mesh.validate()
    return mesh

unit_triangle_vertices = (
    Vector((-3**-0.25, -3**-0.75, 0)),
    Vector((3**-0.25, -3**-0.75, 0)),
    Vector((0, 2/3**0.75, 0)))

def mesh_data_from_matrices(matrices, triangle_scale):
    vertices = []
    polygons = []
    triangle_vertices = [triangle_scale * v for v in unit_triangle_vertices]

    for i, matrix in enumerate(matrices):
        vertices.extend((matrix @ v for v in triangle_vertices))
        polygons.append((i * 3 + 0, i * 3 + 1, i * 3 + 2))

    return vertices, polygons


# Target Provider
#################################################

class BuildState(enum.Enum):
    FINISHED = enum.auto()
    ONGOING = enum.auto()

class TargetProvider:
    def start_build(self, target_object):
        pass

    def continue_build(self, event):
        return BuildState.FINISHED

    def get_matrices(self, scatter_settings):
        return []

    def draw(self):
        pass

class StrokeTarget(TargetProvider):
    def start_build(self, target_object):
        self.points = []
        self.bvhtree = bvhtree_from_object(target_object)
        self.batch = None

    def continue_build(self, event):
        if event.type == 'LEFTMOUSE' and event.value == 'RELEASE':
            return BuildState.FINISHED

        mouse_pos = (event.mouse_region_x, event.mouse_region_y)
        location, *_ = shoot_region_2d_ray(self.bvhtree, mouse_pos)
        if location is not None:
            self.points.append(location)
            self.batch = None
        return BuildState.ONGOING

    def draw(self):
        if self.batch is None:
            self.batch = create_line_strip_batch(self.points)
        draw_line_strip_batch(self.batch, color=(1.0, 0.4, 0.1, 1.0), thickness=5)

    def get_matrices(self, scatter_settings):
        return scatter_around_stroke(self.points, self.bvhtree, scatter_settings)

def scatter_around_stroke(stroke_points, bvhtree, settings):
    scattered_matrices = []
    for point, local_seed in iter_points_on_stroke_with_seed(stroke_points, settings.density, settings.seed):
        matrix = scatter_from_source_point(bvhtree, point, local_seed, settings)
        scattered_matrices.append(matrix)
    return scattered_matrices

def iter_points_on_stroke_with_seed(stroke_points, density, seed):
    for i, (start, end) in enumerate(iter_pairwise(stroke_points)):
        segment_seed = sub_seed(seed, i)
        segment_vector = end - start

        segment_length = segment_vector.length
        amount = round_random(segment_length * density, segment_seed)

        for j in range(amount):
            t = random_uniform(sub_seed(segment_seed, j, 0))
            origin = start + t * segment_vector
            yield origin, sub_seed(segment_seed, j, 1)

def scatter_from_source_point(bvhtree, point, seed, settings):
    # Project displaced point on surface
    radius = random_uniform(sub_seed(seed, 0)) * settings.radius
    offset = random_vector(sub_seed(seed, 2)) * radius
    location, normal, *_ = bvhtree.find_nearest(point + offset)
    assert location is not None
    normal.normalize()

    up_direction = normal if settings.use_normal_rotation else Vector((0, 0, 1))

    # Scale
    min_scale = settings.scale * (1 - settings.random_scale)
    max_scale = settings.scale
    scale = random_uniform(sub_seed(seed, 1), min_scale, max_scale)

    # Location
    location += normal * settings.normal_offset * scale

    # Rotation
    z_rotation = Euler((0, 0, random_uniform(sub_seed(seed, 3), 0, 2 * math.pi))).to_matrix()
    up_rotation = up_direction.to_track_quat('Z', 'X').to_matrix()
    local_rotation = random_euler(sub_seed(seed, 3), settings.rotation).to_matrix()
    rotation = local_rotation @ up_rotation @ z_rotation

    return Matrix.Translation(location) @ rotation.to_4x4() @ scale_matrix(scale)


# Drawing
#################################################

box_vertices = (
    (-1, -1,  1), ( 1, -1,  1), ( 1,  1,  1), (-1,  1,  1),
    (-1, -1, -1), ( 1, -1, -1), ( 1,  1, -1), (-1,  1, -1))

box_indices = (
    (0, 1, 2), (2, 3, 0), (1, 5, 6), (6, 2, 1),
    (7, 6, 5), (5, 4, 7), (4, 0, 3), (3, 7, 4),
    (4, 5, 1), (1, 0, 4), (3, 2, 6), (6, 7, 3))

box_vertices = tuple(Vector(vertex) * 0.5 for vertex in box_vertices)

def draw_matrices_batches(batches):
    shader = get_uniform_color_shader()
    shader.bind()
    shader.uniform_float("color", (0.4, 0.4, 1.0, 0.3))

    bgl.glEnable(bgl.GL_BLEND)
    bgl.glDepthMask(bgl.GL_FALSE)

    for batch in batches:
        batch.draw(shader)

    bgl.glDisable(bgl.GL_BLEND)
    bgl.glDepthMask(bgl.GL_TRUE)

def create_batch_for_matrices(matrices, base_scale):
    coords = []
    indices = []

    scaled_box_vertices = [base_scale * vertex for vertex in box_vertices]

    for matrix in matrices:
        offset = len(coords)
        coords.extend((matrix @ vertex for vertex in scaled_box_vertices))
        indices.extend(tuple(index + offset for index in element) for element in box_indices)

    batch = batch_for_shader(get_uniform_color_shader(),
        'TRIS', {"pos" : coords}, indices = indices)
    return batch


def draw_line_strip_batch(batch, color, thickness=1):
    shader = get_uniform_color_shader()
    bgl.glLineWidth(thickness)
    shader.bind()
    shader.uniform_float("color", color)
    batch.draw(shader)

def create_line_strip_batch(coords):
    return batch_for_shader(get_uniform_color_shader(), 'LINE_STRIP', {"pos" : coords})


def draw_text(location, text, size=15, color=(1, 1, 1, 1)):
    font_id = 0
    ui_scale = bpy.context.preferences.system.ui_scale
    blf.position(font_id, *location)
    blf.size(font_id, round(size * ui_scale), 72)
    blf.draw(font_id, text)


# Utilities
########################################################

'''
Pythons random functions are designed to be used in cases
when a seed is set once and then many random numbers are
generated. To improve the user experience I want to have
full control over how random seeds propagate through the
functions. This is why I use custom random functions.

One benefit is that changing the object density does not
generate new random positions for all objects.
'''

def round_random(value, seed):
    probability = value % 1
    if probability < random_uniform(seed):
        return math.floor(value)
    else:
        return math.ceil(value)

def random_vector(x, min=-1, max=1):
    return Vector((
        random_uniform(sub_seed(x, 0), min, max),
        random_uniform(sub_seed(x, 1), min, max),
        random_uniform(sub_seed(x, 2), min, max)))

def random_euler(x, factor):
    return Euler(tuple(random_vector(x) * factor))

def random_uniform(x, min=0, max=1):
    return random_int(x) / 2147483648 * (max - min) + min

def random_int(x):
    x = (x<<13) ^ x
    return (x * (x * x * 15731 + 789221) + 1376312589) & 0x7fffffff

def sub_seed(seed, index, index2=0):
    return random_int(seed * 3243 + index * 5643 + index2 * 54243)


def currently_in_3d_view(context):
    return context.space_data.type == 'VIEW_3D'

def get_selected_non_active_objects(context):
    return set(context.selected_objects) - {context.active_object}

def make_random_chunks(sequence, chunk_amount):
    sequence = list(sequence)
    random.shuffle(sequence)
    return make_chunks(sequence, chunk_amount)

def make_chunks(sequence, chunk_amount):
    length = math.ceil(len(sequence) / chunk_amount)
    return [sequence[i:i+length] for i in range(0, len(sequence), length)]

def iter_pairwise(sequence):
    return zip(sequence, islice(sequence, 1, None))

def bvhtree_from_object(object):
    import bmesh
    bm = bmesh.new()

    depsgraph = bpy.context.evaluated_depsgraph_get()
    object_eval = object.evaluated_get(depsgraph)
    mesh = object_eval.to_mesh()
    bm.from_mesh(mesh)
    bm.transform(object.matrix_world)

    bvhtree = BVHTree.FromBMesh(bm)
    object_eval.to_mesh_clear()
    return bvhtree

def shoot_region_2d_ray(bvhtree, position_2d):
    region = bpy.context.region
    region_3d = bpy.context.space_data.region_3d

    origin = region_2d_to_origin_3d(region, region_3d, position_2d)
    direction = region_2d_to_vector_3d(region, region_3d, position_2d)

    location, normal, index, distance = bvhtree.ray_cast(origin, direction)
    return location, normal, index, distance

def scale_matrix(factor):
    m = Matrix.Identity(4)
    m[0][0] = factor
    m[1][1] = factor
    m[2][2] = factor
    return m

def event_is_in_region(event, region):
    return (region.x <= event.mouse_x <= region.x + region.width
        and region.y <= event.mouse_y <= region.y + region.height)

def get_max_object_side_length(objects):
    return max(
        max(obj.dimensions[0] for obj in objects),
        max(obj.dimensions[1] for obj in objects),
        max(obj.dimensions[2] for obj in objects)
    )

def get_uniform_color_shader():
    return gpu.shader.from_builtin('3D_UNIFORM_COLOR')


# Registration
###############################################

def register():
    bpy.utils.register_class(ScatterObjects)

def unregister():
    bpy.utils.unregister_class(ScatterObjects)