from math import pi, cos, sin, sqrt
import bpy
from bpy.props import *
from bpy.types import NodeSocket
import mathutils

import arm.utils

if arm.is_reload(__name__):
    arm.utils = arm.reload_module(arm.utils)
else:
    arm.enable_reload(__name__)


def _on_update_socket(self, context):
    self.node.on_socket_val_update(context, self)


class ArmCustomSocket(NodeSocket):
    """
    A custom socket that can be used to define more socket types for
    logic node packs. Do not use this type directly (it is not
    registered)!
    """

    bl_idname = 'ArmCustomSocket'
    bl_label = 'Custom Socket'
    # note: trying to use the `type` property will fail. All custom nodes will have "VALUE" as a type, because it is the default.
    arm_socket_type = 'NONE'
    # please also declare a property named "default_value_raw" of arm_socket_type isn't "NONE"

    def get_default_value(self):
        """Override this for values of unconnected input sockets."""
        return None


class ArmActionSocket(ArmCustomSocket):
    bl_idname = 'ArmNodeSocketAction'
    bl_label = 'Action Socket'
    arm_socket_type = 'NONE'

    def draw(self, context, layout, node, text):
        layout.label(text=self.name)

    def draw_color(self, context, node):
        return 0.8, 0.3, 0.3, 1


class ArmAnimActionSocket(ArmCustomSocket):
    bl_idname = 'ArmNodeSocketAnimAction'
    bl_label = 'Action Socket'
    arm_socket_type = 'STRING'

    default_value_get: PointerProperty(name='Action', type=bpy.types.Action)  # legacy version of the line after this one
    default_value_raw: PointerProperty(name='Action', type=bpy.types.Action, update=_on_update_socket)

    def __init__(self):
        super().__init__()
        if self.default_value_get is not None:
            self.default_value_raw = self.default_value_get
            self.default_value_get = None

    def get_default_value(self):
        if self.default_value_raw is None:
            return ''
        if self.default_value_raw.name not in bpy.data.actions:
            return self.default_value_raw.name
        name = arm.utils.asset_name(bpy.data.actions[self.default_value_raw.name])
        return arm.utils.safestr(name)

    def draw(self, context, layout, node, text):
        if self.is_output:
            layout.label(text=self.name)
        elif self.is_linked:
            layout.label(text=self.name)
        else:
            layout.prop_search(self, 'default_value_raw', bpy.data, 'actions', icon='NONE', text='')

    def draw_color(self, context, node):
        return 0.8, 0.8, 0.8, 1

    
class ArmRotationSocket(ArmCustomSocket):
    bl_idname = 'ArmRotationSocket'
    bl_label = 'Rotation Socket'
    arm_socket_type = 'ROTATION'  # the internal representation is a quaternion, AKA a '4D vector' (using mathutils.Vector((x,y,z,w)))
    
    def get_default_value(self):
        if self.default_value_raw is None:
            return mathutils.Vector((0.0,0.0,0.0,1.0))
        else:
            return self.default_value_raw

    def on_unit_update(self, context):
        if self.default_value_unit == 'Rad':
            fac = pi/180  # deg->rad conversion
        else:
            fac = 180/pi  # rad->deg conversion
        if self.default_value_mode == 'AxisAngle':
            self.default_value_s3 *= fac
        elif self.default_value_mode == 'EulerAngles':
            self.default_value_s0 *= fac
            self.default_value_s1 *= fac
            self.default_value_s2 *= fac
        self.do_update_raw(context)
        
            
    def on_mode_update(self, context):
        if self.default_value_mode == 'Quaternion':
            summ = abs(self.default_value_s0)
            summ+= abs(self.default_value_s1)
            summ+= abs(self.default_value_s2)
            summ+= abs(self.default_value_s3)
            if summ<0.01:
                self.default_value_s3 = 1.0
        elif self.default_value_mode == 'AxisAngle':
            summ = abs(self.default_value_s0)
            summ+= abs(self.default_value_s1)
            summ+= abs(self.default_value_s2)
            if summ<1E-5:
                self.default_value_s3 = 0.0
        self.do_update_raw(context)

    @staticmethod
    def convert_to_quaternion(part1,part2,param1,param2,param3):
        """converts a representation of rotation into a quaternion.
        ``part1`` is a vector, ``part2`` is a scalar or None,
        ``param1`` is in ('Quaternion', 'EulerAngles', 'AxisAngle'),
        ``param2`` is in ('Rad','Deg') for both EulerAngles and AxisAngle,
        ``param3`` is a len-3 string like "XYZ", for EulerAngles """
        if param1=='Quaternion':
            qx, qy, qz = part1[0], part1[1], part1[2]
            qw = part2
            # need to normalize the quaternion for a rotation (having it be 0 is not an option)
            ql = sqrt(qx**2+qy**2+qz**2+qw**2)
            if abs(ql)<1E-5:
                qx, qy, qz, qw = 0.0,0.0,0.0,1.0
            else:
                qx /= ql
                qy /= ql
                qz /= ql
                qw /= ql
            return mathutils.Vector((qx,qy,qz,qw))
        
        elif param1 == 'AxisAngle':
            if param2 == 'Deg':
                angle = part2 * pi/180
            else:
                angle = part2
            cang, sang = cos(angle/2), sin(angle/2)
            x,y,z = part1[0], part1[1], part1[2]
            veclen = sqrt(x**2+y**2+z**2)
            if veclen<1E-5:
                return mathutils.Vector((0.0,0.0,0.0,1.0))
            else:
                return mathutils.Vector((
                    x/veclen * sang,
                    y/veclen * sang,
                    z/veclen * sang,
                    cang
                ))
        else:  # param1 == 'EulerAngles'
            x,y,z = part1[0], part1[1], part1[2]
            if param2 == 'Deg':
                x *= pi/180
                y *= pi/180
                z *= pi/180
            cx, sx = cos(x/2), sin(x/2)
            cy, sy = cos(y/2), sin(y/2)
            cz, sz = cos(z/2), sin(z/2)

            qw, qx, qy, qz  = 1.0,0.0,0.0,0.0
            for direction in param3[::-1]:
                qwi, qxi,qyi,qzi = {'X': (cx,sx,0,0), 'Y': (cy,0,sy,0), 'Z': (cz,0,0,sz)}[direction]

                qw = qw*qwi -qx*qxi -qy*qyi -qz*qzi
                qx = qx*qwi +qw*qxi +qy*qzi -qz*qyi
                qy = qy*qwi +qw*qyi +qz*qxi -qx*qzi
                qz = qz*qwi +qw*qzi +qx*qyi -qy*qxi
            return mathutils.Vector((qx,qy,qz,qw))

        
    def do_update_raw(self, context):
        part1 = mathutils.Vector((
            self.default_value_s0,
            self.default_value_s1,
            self.default_value_s2, 1
        ))
        part2 = self.default_value_s3

        self.default_value_raw = self.convert_to_quaternion(
            part1,
            self.default_value_s3,
            self.default_value_mode,
            self.default_value_unit,
            self.default_value_order
        )
            
        
    def draw(self, context, layout, node, text):
        if (self.is_output or self.is_linked):
            layout.label(text=self.name)
        else:
            coll1 = layout.column(align=True)
            coll1.label(text=self.name)
            bx=coll1.box()
            coll = bx.column(align=True)
            coll.prop(self, 'default_value_mode')
            if self.default_value_mode in ('EulerAngles', 'AxisAngle'):
                coll.prop(self, 'default_value_unit')
                
            if self.default_value_mode == 'EulerAngles':
                coll.prop(self, 'default_value_order')
                coll.prop(self, 'default_value_s0', text='X')
                coll.prop(self, 'default_value_s1', text='Y')
                coll.prop(self, 'default_value_s2', text='Z')
            elif self.default_value_mode == 'Quaternion':
                coll.prop(self, 'default_value_s0', text='X')
                coll.prop(self, 'default_value_s1', text='Y')
                coll.prop(self, 'default_value_s2', text='Z')
                coll.prop(self, 'default_value_s3', text='W')
            elif self.default_value_mode == 'AxisAngle':
                coll.prop(self, 'default_value_s0', text='X')
                coll.prop(self, 'default_value_s1', text='Y')
                coll.prop(self, 'default_value_s2', text='Z')
                coll.separator()
                coll.prop(self, 'default_value_s3', text='Angle')            

    def draw_color(self, context, node):
        return 0.68, 0.22, 0.62, 1

    default_value_mode: EnumProperty(
        items=[('EulerAngles', 'Euler Angles', 'Euler Angles'),
               ('AxisAngle', 'Axis/Angle', 'Axis/Angle'),
               ('Quaternion', 'Quaternion', 'Quaternion')],
        name='', default='EulerAngles',
        update=on_mode_update)

    default_value_unit: EnumProperty(
        items=[('Deg', 'Degrees', 'Degrees'),
               ('Rad', 'Radians', 'Radians')],
        name='', default='Rad',
        update=on_unit_update)
    default_value_order: EnumProperty(
        items=[('XYZ','XYZ','XYZ'),
               ('XZY','XZY (legacy Armory euler order)','XZY (legacy Armory euler order)'),
               ('YXZ','YXZ','YXZ'),
               ('YZX','YZX','YZX'),
               ('ZXY','ZXY','ZXY'),
               ('ZYX','ZYX','ZYX')],
        name='', default='XYZ'
    )

    default_value_s0: FloatProperty(update=do_update_raw)
    default_value_s1: FloatProperty(update=do_update_raw)
    default_value_s2: FloatProperty(update=do_update_raw)
    default_value_s3: FloatProperty(update=do_update_raw)

    default_value_raw: FloatVectorProperty(
        name='Value',
        description='Raw quaternion obtained for the default value of a ArmRotationSocket socket',
        size=4, default=(0,0,0,1),
        update = _on_update_socket
    )


class ArmArraySocket(ArmCustomSocket):
    bl_idname = 'ArmNodeSocketArray'
    bl_label = 'Array Socket'
    arm_socket_type = 'NONE'

    def draw(self, context, layout, node, text):
        layout.label(text=self.name)

    def draw_color(self, context, node):
        return 0.8, 0.4, 0.0, 1


class ArmBoolSocket(ArmCustomSocket):
    bl_idname = 'ArmBoolSocket'
    bl_label = 'Boolean Socket'
    arm_socket_type = 'BOOLEAN'

    default_value_raw: BoolProperty(
        name='Value',
        description='Input value used for unconnected socket',
        update=_on_update_socket
    )

    def draw(self, context, layout, node, text):
        draw_socket_layout(self, layout)

    def draw_color(self, context, node):
        return 0.8, 0.651, 0.839, 1

    def get_default_value(self):
        return self.default_value_raw


class ArmColorSocket(ArmCustomSocket):
    bl_idname = 'ArmColorSocket'
    bl_label = 'Color Socket'
    arm_socket_type = 'RGBA'

    default_value_raw: FloatVectorProperty(
        name='Value',
        size=4,
        subtype='COLOR',
        min=0.0,
        max=1.0,
        default=[0.0, 0.0, 0.0, 1.0],
        description='Input value used for unconnected socket',
        update=_on_update_socket
    )

    def draw(self, context, layout, node, text):
        draw_socket_layout_split(self, layout)

    def draw_color(self, context, node):
        return 0.78, 0.78, 0.161, 1

    def get_default_value(self):
        return self.default_value_raw


class ArmDynamicSocket(ArmCustomSocket):
    bl_idname = 'ArmDynamicSocket'
    bl_label = 'Dynamic Socket'
    arm_socket_type = 'NONE'

    def draw(self, context, layout, node, text):
        layout.label(text=self.name)

    def draw_color(self, context, node):
        return 0.388, 0.78, 0.388, 1


class ArmFloatSocket(ArmCustomSocket):
    bl_idname = 'ArmFloatSocket'
    bl_label = 'Float Socket'
    arm_socket_type = 'VALUE'

    default_value_raw: FloatProperty(
        name='Value',
        description='Input value used for unconnected socket',
        precision=3,
        update=_on_update_socket
    )

    def draw(self, context, layout, node, text):
        draw_socket_layout(self, layout)

    def draw_color(self, context, node):
        return 0.631, 0.631, 0.631, 1

    def get_default_value(self):
        return self.default_value_raw


class ArmIntSocket(ArmCustomSocket):
    bl_idname = 'ArmIntSocket'
    bl_label = 'Integer Socket'
    arm_socket_type = 'INT'

    default_value_raw: IntProperty(
        name='Value',
        description='Input value used for unconnected socket',
        update=_on_update_socket
    )

    def draw(self, context, layout, node, text):
        draw_socket_layout(self, layout)

    def draw_color(self, context, node):
        return 0.059, 0.522, 0.149, 1

    def get_default_value(self):
        return self.default_value_raw


class ArmObjectSocket(ArmCustomSocket):
    bl_idname = 'ArmNodeSocketObject'
    bl_label = 'Object Socket'
    arm_socket_type = 'OBJECT'

    default_value_get: PointerProperty(name='Object', type=bpy.types.Object)  # legacy version of the line after this one
    default_value_raw: PointerProperty(name='Object', type=bpy.types.Object, update=_on_update_socket)

    def __init__(self):
        super().__init__()
        if self.default_value_get is not None:
            self.default_value_raw = self.default_value_get
            self.default_value_get = None

    def get_default_value(self):
        if self.default_value_raw is None:
            return ''
        if self.default_value_raw.name not in bpy.data.objects:
            return self.default_value_raw.name
        return arm.utils.asset_name(bpy.data.objects[self.default_value_raw.name])

    def draw(self, context, layout, node, text):
        if self.is_output:
            layout.label(text=self.name)
        elif self.is_linked:
            layout.label(text=self.name)
        else:
            row = layout.row(align=True)
            row.prop_search(self, 'default_value_raw', context.scene, 'objects', icon='NONE', text=self.name)

    def draw_color(self, context, node):
        return 0.15, 0.55, 0.75, 1



class ArmStringSocket(ArmCustomSocket):
    bl_idname = 'ArmStringSocket'
    bl_label = 'String Socket'
    arm_socket_type = 'STRING'

    default_value_raw: StringProperty(
        name='Value',
        description='Input value used for unconnected socket',
        update=_on_update_socket
    )

    def draw(self, context, layout, node, text):
        draw_socket_layout_split(self, layout)

    def draw_color(self, context, node):
        return 0.439, 0.698, 1, 1

    def get_default_value(self):
        return self.default_value_raw


class ArmVectorSocket(ArmCustomSocket):
    bl_idname = 'ArmVectorSocket'
    bl_label = 'Vector Socket'
    arm_socket_type = 'VECTOR'

    default_value_raw: FloatVectorProperty(
        name='Value',
        size=3,
        precision=3,
        description='Input value used for unconnected socket',
        update=_on_update_socket
    )

    def draw(self, context, layout, node, text):
        if not self.is_output and not self.is_linked:
            col = layout.column(align=True)
            col.label(text=self.name + ":")
            col.prop(self, 'default_value_raw', text='')
        else:
            layout.label(text=self.name)

    def draw_color(self, context, node):
        return 0.388, 0.388, 0.78, 1

    def get_default_value(self):
        return self.default_value_raw


def draw_socket_layout(socket: bpy.types.NodeSocket, layout: bpy.types.UILayout, prop_name='default_value_raw'):
    if not socket.is_output and not socket.is_linked:
        layout.prop(socket, prop_name, text=socket.name)
    else:
        layout.label(text=socket.name)


def draw_socket_layout_split(socket: bpy.types.NodeSocket, layout: bpy.types.UILayout, prop_name='default_value_raw'):
    if not socket.is_output and not socket.is_linked:
        # Blender layouts use 0.4 splits
        layout = layout.split(factor=0.4, align=True)

    layout.label(text=socket.name)

    if not socket.is_output and not socket.is_linked:
        layout.prop(socket, prop_name, text='')

REG_CLASSES = (
    ArmActionSocket,
    ArmAnimActionSocket,
    ArmRotationSocket,
    ArmArraySocket,
    ArmBoolSocket,
    ArmColorSocket,
    ArmDynamicSocket,
    ArmFloatSocket,
    ArmIntSocket,
    ArmObjectSocket,
    ArmStringSocket,
    ArmVectorSocket,
)
register, unregister = bpy.utils.register_classes_factory(REG_CLASSES)
