# ##### 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 bpy
import os
import sys
import platform
from ctypes import *

CFloat2 = c_float * 2
CFloat3 = c_float * 3
CDouble3 = c_double * 3


class EditArgs(Structure):
    _fields_ = [
        ("brush_type", c_short),
        ("widget_unit", c_short),
        ("smooth_mode", c_short),
        ("orient_mode", c_short),
        ("attract_mode", c_short),
        ("weight_mode", c_short),
        ("continuous_attract", c_short),
        ("attract_method", c_short),
        ("falloff_shape", c_short),
        ("flag", c_int),
        ("noise_flag", c_int),
        ("rad_3d", c_float),
        ("strength_adjust", c_float),
        ("limited_length", c_float),
        ("curving", c_float),
        ("random_control", c_float),
        ("percent_control", c_float),
        ("head_strength", c_float),
        ("middle_strength", c_float),
        ("rough_seed", c_int),
        ("clone_mode", c_short),
        ("clone_coordinate", c_short),
        ("clone_rotation_angle", c_float),
        ("clip_range", CFloat2),
        ("brush_color", CDouble3),
        ("brush_data", c_void_p)
    ]


class ApplyArgs(Structure):
    _fields_ = [
        ("edit_data", c_void_p),
        ("mouse", CFloat2),
        ("mouse_3d", CFloat3),
        ("mouse_nor_3d", CFloat3),
        ("flip", c_bool),
    ]


class CloneDraw(Structure):
    _fields_ = [
        ("pyscene", c_void_p),
        ("pyob", c_void_p),
        ("pypsys", c_void_p),
        ("mouse_3d", CFloat3),
        ("mouse_nor_3d", CFloat3),
        ("rad_3d", c_float),
        ("coord_mode", c_short),
        ("rotate_angle", c_float),
        ("r_cos", c_void_p),
        ("r_cos_len", c_int),
        ("r_indices", c_void_p),
        ("r_indices_len", c_int),
    ]


def init_module():
    lib_name = "hair_brush_3d"
    version = bpy.app.version
    if version[:2] == (3, 2):
        lib_name += "_320_py310"
    elif version[:2] == (3, 1):
        lib_name += "_310_py310"
    elif version[:2] == (3, 0):
        lib_name += "_300_py39"
    elif version[:2] == (2, 93):
        if bpy.app.version[2] < 4 or (bpy.app.version[2] == 4 and bpy.app.version_cycle != 'release'):
            lib_name += "_293_py39"
        else:
            lib_name += "_2934_py39"

    dirname = os.path.dirname(__file__)

    if sys.platform.startswith("linux"):
        ext = ".so"
    else:
        assert sys.platform.startswith("darwin")
        ext = ".dylib"
        if platform.machine() == 'arm64':
            dirname += "/darwin_arm64"

    dylib = cdll.LoadLibrary(f'{dirname}/{lib_name}{ext}')
    dylib.set_mesh_surface_gpu_batch.argtypes = (c_void_p, c_void_p)
    dylib.set_mesh_surface_gpu_batch.restype = c_void_p

    dylib.restore_gpu_batch.argtypes = (c_void_p, c_void_p)

    dylib.init_brush_edit_data.argtypes = (c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p)
    dylib.init_brush_edit_data.restype = c_void_p

    dylib.apply_brush_edit.argtypes = (c_void_p,)
    dylib.apply_brush_edit.restype = c_bool

    dylib.free_brush_edit_data.argtypes = (c_void_p,)

    dylib.set_clone_source_wrap.argtypes = (c_void_p,)
    dylib.set_clone_source_wrap.restype = c_int

    dylib.get_clone_source_for_draw_wrap.argtypes = (c_void_p,)

    dylib.set_hair_color_group_wrap.argtypes = (c_void_p, c_void_p, c_int)

    dylib.get_hair_color_groups.argtypes = (c_void_p, c_void_p)
    dylib.get_hair_color_groups.restype = c_void_p

    dylib.extend_color_array_wrap.argtypes = (c_void_p, c_int)

    dylib.shrink_color_array_wrap.argtypes = (c_void_p, c_void_p, c_int)

    dylib.get_hair_color_wrap.argtypes = (c_void_p, c_void_p)
    dylib.get_hair_color_wrap.restype = c_bool

    dylib.imbuf_to_float_buffer.argtypes = (c_void_p, c_void_p)
    dylib.imbuf_to_float_buffer.restype = c_bool

    dylib.geom_icon_to_image.argtypes = (c_void_p, c_void_p, c_int, c_void_p)
    dylib.geom_icon_to_image.restype = c_bool

    dylib.set_particle_edit_tool_update_func.argtypes = (c_void_p, c_int)

    dylib.set_mesh_calc_normals_function.argtypes = (c_void_p,)

    dylib.get_widget_unit.argtypes = (c_void_p,)
    dylib.get_widget_unit.restype = c_short

    dylib.convert_curve_to_hair.argtypes = (c_void_p, c_void_p, c_void_p)

    dylib.free_c_memory.argtypes = (c_void_p,)
    return dylib


py_clib = init_module()


def set_mesh_surface_gpu_batch(gpu_batch, mesh):
    return py_clib.set_mesh_surface_gpu_batch(id(gpu_batch), mesh)


def restore_gpu_batch(gpu_batch, ori_batch):
    return py_clib.restore_gpu_batch(id(gpu_batch), ori_batch)


def init_brush_edit_data(pydeps, pyscene, pyob, pypsys, pyregion, pyv3d, pybuf, args):
    edit_args = EditArgs()
    for k, v in args.items():
        if k == "clip_range":
            edit_args.__setattr__(k, CFloat2(*v))
        elif k == "brush_color":
            edit_args.__setattr__(k, CDouble3(*v))
        else:
            edit_args.__setattr__(k, v)

    return py_clib.init_brush_edit_data(pydeps, pyscene, pyob, pypsys, pyregion, pyv3d, id(pybuf), byref(edit_args))


def apply_brush_edit(pyedit_data, pymouse, pymouse_3d, pymouse_nor_3d, flip):
    args = ApplyArgs()
    args.edit_data = pyedit_data
    args.mouse = CFloat2(*pymouse)
    args.mouse_3d = CFloat3(*pymouse_3d)
    args.mouse_nor_3d = CFloat3(*pymouse_nor_3d)
    args.flip = flip

    return py_clib.apply_brush_edit(byref(args))


def free_brush_edit_data(edit_data):
    return py_clib.free_brush_edit_data(edit_data)


def set_clone_source(pyedit_data, pymouse, pymouse_3d, pymouse_nor_3d):
    args = ApplyArgs()
    args.edit_data = pyedit_data
    args.mouse = CFloat2(*pymouse)
    args.mouse_3d = CFloat3(*pymouse_3d)
    args.mouse_nor_3d = CFloat3(*pymouse_nor_3d)
    return py_clib.set_clone_source_wrap(byref(args))


def get_clone_source_for_draw(pyscene, pyob, pypsys, mouse_3d, mouse_nor_3d, rad_3d, coord_mode, rotate_angle):
    args = CloneDraw()
    args.pyscene = pyscene
    args.pyob = pyob
    args.pypsys = pypsys
    args.mouse_3d = mouse_3d
    args.mouse_nor_3d = mouse_nor_3d
    args.rad_3d = rad_3d
    args.coord_mode = coord_mode
    args.rotate_angle = rotate_angle

    py_clib.get_clone_source_for_draw_wrap(byref(args))
    if not args.r_cos or not args.r_indices:
        return
    r_cos = cast(args.r_cos, POINTER(c_float))
    r_index = cast(args.r_indices, POINTER(c_int))
    cos = []
    indices = []
    for i in range(0, args.r_cos_len, 3):
        cos.append((r_cos[i], r_cos[i + 1], r_cos[i + 2]))
    for i in range(0, args.r_indices_len, 2):
        indices.append((r_index[i], r_index[i + 1]))

    py_clib.free_c_memory(args.r_cos)
    py_clib.free_c_memory(args.r_indices)
    return cos, indices


def reset_clone_source():
    return py_clib.reset_clone_source_wrap()


def set_hair_color_group(pypsys, pycolor, operate_code):
    return py_clib.set_hair_color_group_wrap(pypsys, byref(CDouble3(*pycolor)), operate_code)


def get_hair_color_groups(psys):
    length = c_int(0)
    groups_ptr = py_clib.get_hair_color_groups(psys, byref(length))
    groups = cast(groups_ptr, POINTER(c_double))
    colors = []
    for i in range(length.value):
        colors.append(groups[i])

    return colors


def extend_color_array(pypsys, length_add):
    return py_clib.extend_color_array_wrap(pypsys, length_add)


def shrink_color_array(pypsys, py_index_list):
    length = len(py_index_list)
    if length == 0:
        return
    index_list = (c_int * length)(*py_index_list)

    return py_clib.shrink_color_array_wrap(pypsys, byref(index_list), length)


def get_hair_color(pypsys, pybuf):
    return py_clib.get_hair_color_wrap(pypsys, id(pybuf))


def imbuf_to_float_buffer(pyibuf, pybuf):
    return py_clib.imbuf_to_float_buffer(id(pyibuf), id(pybuf))


def geom_icon_to_image(py_coords, py_colors, coords_len, pyibuf):
    coords = (c_uint8 * len(py_coords))(*py_coords)
    colors = (c_uint8 * len(py_colors))(*py_colors)

    return py_clib.geom_icon_to_image(byref(coords), byref(colors), coords_len, id(pyibuf))


def set_particle_edit_tool_update_func(prop, do_set):
    return py_clib.set_particle_edit_tool_update_func(prop, do_set)


def set_mesh_calc_normals_function(func):
    return py_clib.set_mesh_calc_normals_function(func)


def get_widget_unit(preferences):
    return py_clib.get_widget_unit(preferences)


def convert_curve_to_hair(mesh_obj, psys, curve_obj):
    return py_clib.convert_curve_to_hair(mesh_obj, psys, curve_obj)
