from math import pi

import bpy
from bpy.types import Context, Operator


def create_triplanar_setup(context: Context) -> None:
    """
    Creates triplanar setup - a mesh cube with UVProject modifier and 6 UV projectors

    :param context: Current context
    :type context: bpy.types.Context
    """
    bpy.ops.object.empty_add(type='CUBE', radius=1.6, location=(1.6, 1.6, 1.6), rotation=(0, 0, 0))
    anchor_empty = context.active_object
    anchor_empty.name = "UV_anchor"

    bpy.ops.mesh.primitive_cube_add(radius=1, calc_uvs=True)
    cube_obj = context.active_object

    bpy.ops.object.modifier_add(type='UV_PROJECT')
    uv_mod = cube_obj.modifiers["UVProject"]

    rotations_list = [
        ("top", (0, 0, 0)),
        ("front", (pi/2, 0, 0)),
        ("back", (pi/2, 0, pi)),
        ("left", (pi/2, 0, pi/2)),
        ("right", (pi/2, 0, -pi/2)),
        ("bottom", (pi, 0, 0)),
    ]

    uv_mod.projector_count = len(rotations_list)

    for i, (name, rotation) in enumerate(rotations_list):
        bpy.ops.object.empty_add(type='SINGLE_ARROW', radius=1.6, location=(0, 0, 0), rotation=rotation)
        uv_projector = context.active_object
        uv_projector.name = "UV_%s" % name
        uv_projector.parent = anchor_empty
        uv_mod.projectors[i].object = uv_projector


class MESH_OT_AddTriplanarSetup(Operator):
    bl_idname = "mesh.setup_triplanar_add"
    bl_label = "Triplanar UV Setup"
    bl_options = {'REGISTER', 'UNDO'}

    def execute(self, context: Context):
        try:
            create_triplanar_setup(context)
        except Exception as e:
            self.report({'ERROR'}, str(e))
            return {'CANCELLED'}
        return {'FINISHED'}


def draw_triplanar_setup(self, context):
    self.layout.operator(MESH_OT_AddTriplanarSetup.bl_idname, icon='MOD_UVPROJECT')


def register():
    bpy.utils.register_class(MESH_OT_AddTriplanarSetup)
    if draw_triplanar_setup.__name__ not in [f.__name__ for f in bpy.types.INFO_MT_mesh_add._dyn_ui_initialize()]:
        bpy.types.INFO_MT_mesh_add.append(draw_triplanar_setup)


def unregister():
    bpy.types.INFO_MT_mesh_add.remove(draw_triplanar_setup)
    bpy.utils.unregister_class(MESH_OT_AddTriplanarSetup)


register()