# Copyright (C) 2022 Thomas Hoppe (h0bB1T). All rights reserved.
#
# Unauthorized copying of this file via any medium is strictly prohibited.
# Proprietary and confidential.

import bpy

from dataclasses import dataclass

from bpy.types import Menu, UILayout

from ..utils.tools import local_addon_info
from ..properties import Properties
from ..utils.blender import asset_wizard_pro_enabled, is_ls_410
from ..utils.colors import colors
from ..registries.collections import Collections
from ..registries.icons_registry import IconsRegistry
from ..operators.create_node_std import UI_OT_create_node_std
from ..operators.create_node_new_image import UI_OT_create_node_new_image
from ..operators.create_node_uv_map import UI_OT_create_node_uv_map
from ..operators.create_node_vertex_colors import UI_OT_create_node_vertex_colors
from ..operators.delete_element import UI_OT_delete_element
from ..registries.texture_registry import RepositoryType, TextureRegistry
from ..menus.node_grunge_pie_menus import NODE_MT_snw_grunge_category_pie_menu
from ..menus.node_image_pie_menus import NODE_MT_snw_image_category_pie_menu
from ..menus.node_pbr_pie_menus import NODE_MT_snw_pbr_category_pie_menu


@dataclass
class ShdStd:
    text: str
    icon: str
    node: str
    operation: str = None
    blend_type: str = None


class NODE_MT_snw_main_pie_menu(Menu):
    """
    The primary PIE menu, shown when hitting the primary addon button ('D').
    """
    bl_idname = "NODE_MT_snw_main_pie_menu"
    bl_label = "Shader Node Wizard"


    @classmethod
    def poll(cls, context):
        return context.space_data.type == 'NODE_EDITOR' and \
            context.space_data.node_tree is not None  and \
            context.space_data.tree_type == 'ShaderNodeTree'  
      

    def __draw_tools(self, layout: UILayout):
        c = layout
        c.ui_units_x = 30

        # Check presence of SNW grunge set
        if TextureRegistry.instance().get_by_bin_hash('980273b26258d937044e39c69302a5d61efdecaf8f2acd44ae5666f49c37ad60@0')[1]:
            visible_columns = 7
            g = c.grid_flow(row_major=True, columns=visible_columns, even_columns=True, align=True)
            for col in Collections.instance().static_collections():
                if col.display == 'menu':
                    menus = {}
                    for idx, it in enumerate(col.items):
                        if it.menu_id not in menus:
                            menus[it.menu_id] = [] 
                        menus[it.menu_id].append((idx, it))

                    for k in sorted(menus.keys()):
                        b = g.column()
                        for idx, it in menus[k]:
                            op = b.operator('snw.create_ci_node', text=it.name, icon=it.op_icon)
                            op.group = it.group
                            op.id = idx
                            op.description = it.description

                    for it in range(visible_columns - len(menus)):
                        g.label(text='')
                else:
                    for idx, it in enumerate(col.items):
                        b = g.box().column(align=True)
                        b.template_icon(it.icon_id, scale=5)
                        op = b.operator('snw.create_ci_node', text=it.name, icon=it.op_icon)
                        op.group = it.group
                        op.id = idx
                        op.description = it.description
                    for it in range(visible_columns - (len(col.items) % visible_columns)):
                        g.label(text='')
        else:
            c.box().label(text='Missing SNW Grunge Textures, required for Collection Nodes.')


    def __std_simple_shader(self, l: UILayout, node: ShdStd, color: str = None):
        UI_OT_create_node_std.create_ui(
            l, 
            'Shader', 
            node.node, 
            f"'{node.operation}'" if node.operation else '', 
            f"'{node.blend_type}'" if node.blend_type else '',
            node.text, 
            node.icon, 
            f'Drop a "{node.text}" node', 
            color=color
        )


    def __draw_std_nodes_tools(self, l: UILayout):
        nodes = {
            'Utils': [
                ShdStd( 'Map Range', 'UV_SYNC_SELECT', 'MapRange' ),
                ShdStd( 'Combine XYZ', 'FULLSCREEN_EXIT', 'CombineXYZ' ),
                ShdStd( 'Separate XYZ', 'FULLSCREEN_ENTER', 'SeparateXYZ' ),
                ShdStd( 'Blackbody', 'LIGHT_SUN', 'Blackbody' ),
            ],
            'Info': [
                ShdStd( 'Geometry', 'MOD_WIREFRAME', 'NewGeometry' ),
                ShdStd( 'Texture Coordinate', 'UV', 'TexCoord' ),
                ShdStd( 'Layer Weight', 'MATSPHERE', 'LayerWeight' ),
                ShdStd( 'Light Path', 'LIGHT_SPOT', 'LightPath' ),
                ShdStd( 'Ambient Occlusion', 'OVERLAY', 'AmbientOcclusion' ),
            ],
            'Texture': [
                ShdStd( 'Noise Texture', 'MOD_NOISE', 'TexNoise' ),
            ] 
            + ([ShdStd( 'Musgrave Texture', 'RNDCURVE', 'TexMusgrave' ),] if is_ls_410() else [])
            + [
                ShdStd( 'Voronoi Texture', 'LIGHTPROBE_PLANAR', 'TexVoronoi' ),
                ShdStd( 'White Noise', 'SEQ_HISTOGRAM', 'TexWhiteNoise' ),
                ShdStd( 'Image Texture', 'IMAGE_DATA', 'TexImage' ),
            ],
            'Color': [
                ShdStd( 'Hue Saturation', 'COLOR', 'HueSaturation' ),
                ShdStd( 'Brightness Contrast', 'IMAGE_ZDEPTH', 'BrightContrast' ),
                ShdStd( 'ColorRamp', 'NODE_TEXTURE', 'ValToRGB' ),
                ShdStd( 'Gamma', 'IPO_EASE_OUT', 'Gamma' ),
                ShdStd( 'Invert', 'OVERLAY', 'Invert' ),
            ],
            'Normal': [
                ShdStd( 'Bump Map', 'FORCE_TURBULENCE', 'Bump' ),
                ShdStd( 'Normal Map', 'MOD_OCEAN', 'NormalMap' ),
            ],
            'Shader': [
                ShdStd( 'Principled BSDF', 'MATERIAL', 'BsdfPrincipled' ),
                ShdStd( 'Mix Shader', 'ORIENTATION_GIMBAL', 'MixShader' ),
                ShdStd( 'Emission', 'LIGHT', 'Emission' ),
                ShdStd( 'Transparent', 'MOD_UVPROJECT', 'BsdfTransparent' ),
                ShdStd( 'Principled Volume', 'OUTLINER_DATA_VOLUME', 'VolumePrincipled' ),
                ShdStd( 'Glass BSDF', 'NODE_MATERIAL', 'BsdfGlass' ),
            ],
            'Input': [
                ShdStd( 'Attribute', 'NODETREE', 'Attribute' ),
                ShdStd( 'RGB', 'COLOR', 'RGB' ),
                ShdStd( 'Value', 'ONIONSKIN_ON', 'Value' ),
            ],
            'Math': [
                ShdStd( 'Add', 'DRIVER_TRANSFORM', 'Math', operation='ADD' ),
                ShdStd( 'Subtract', 'DRIVER_TRANSFORM', 'Math', operation='SUBTRACT' ),
                ShdStd( 'Multiply', 'DRIVER_TRANSFORM', 'Math', operation='MULTIPLY'),
                ShdStd( 'Divide', 'DRIVER_TRANSFORM', 'Math', operation='DIVIDE'),
                ShdStd( 'Multiply Add', 'DRIVER_TRANSFORM', 'Math', operation='MULTIPLY_ADD'),
                (),
                ShdStd( 'Power', 'DRIVER_TRANSFORM', 'Math', operation='POWER'),
                ShdStd( 'Absolute', 'DRIVER_TRANSFORM', 'Math', operation='ABSOLUTE'),
                (),
                ShdStd( 'Minimum', 'DRIVER_TRANSFORM', 'Math', operation='MINIMUM'),
                ShdStd( 'Maximum', 'DRIVER_TRANSFORM', 'Math', operation='MAXIMUM'),
                ShdStd( 'Less Than', 'DRIVER_TRANSFORM', 'Math', operation='LESS_THAN'),
                ShdStd( 'Greater Than', 'DRIVER_TRANSFORM', 'Math', operation='GREATER_THAN'),
                ShdStd( 'Smooth Min', 'DRIVER_TRANSFORM', 'Math', operation='SMOOTH_MIN'),
                ShdStd( 'Smooth Max', 'DRIVER_TRANSFORM', 'Math', operation='SMOOTH_MAX'),
                (),
                ShdStd( 'Round', 'DRIVER_TRANSFORM', 'Math', operation='ROUND'),
                ShdStd( 'Floor', 'DRIVER_TRANSFORM', 'Math', operation='FLOOR'),
                ShdStd( 'Ceil', 'DRIVER_TRANSFORM', 'Math', operation='CEIL'),
                (),
                ShdStd( 'Fraction', 'DRIVER_TRANSFORM', 'Math', operation='FRACT'),
                ShdStd( 'Modulo', 'DRIVER_TRANSFORM', 'Math', operation='MODULO'),
                ShdStd( 'Warp', 'DRIVER_TRANSFORM', 'Math', operation='WRAP'),
                ShdStd( 'Snap', 'DRIVER_TRANSFORM', 'Math', operation='SNAP'),
                ShdStd( 'Ping Pong', 'DRIVER_TRANSFORM', 'Math', operation='PINGPONG'),
                (),
                ShdStd( 'Sin', 'DRIVER_TRANSFORM', 'Math', operation='SINE'),
                ShdStd( 'Cos', 'DRIVER_TRANSFORM', 'Math', operation='COSINE'),
                ShdStd( 'Tan', 'DRIVER_TRANSFORM', 'Math', operation='TANGENT'),
            ],
            'VMath': [
                ShdStd( 'Add', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='ADD' ),
                ShdStd( 'Subtract', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='SUBTRACT'),
                ShdStd( 'Multiply', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='MULTIPLY'),
                ShdStd( 'Divide', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='DIVIDE'),
                ShdStd( 'Multiply Add', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='MULTIPLY_ADD'),
                (),
                ShdStd( 'Cross', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='CROSS_PRODUCT'),
                ShdStd( 'Project', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='PROJECT'),
                ShdStd( 'Reflect', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='REFLECT'),
                ShdStd( 'Refract', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='REFRACT'),
                ShdStd( 'Faceforward', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='FACEFORWARD'),
                ShdStd( 'Dot', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='DOT_PRODUCT'),
                (),
                ShdStd( 'Distance', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='DISTANCE'),
                ShdStd( 'Lenth', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='LENGTH'),
                ShdStd( 'Scale', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='SCALE'),
                ShdStd( 'Normalize', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='NORMALIZE'),
                (),
                ShdStd( 'Absolute', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='ABSOLUTE'),
                ShdStd( 'Minimum', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='MINIMUM'),
                ShdStd( 'Maximum', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='MAXIMUM'),
                ShdStd( 'Floor', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='FLOOR'),
                ShdStd( 'Ceil', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='CEIL'),
                ShdStd( 'Fraction', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='FRACTION'),
                ShdStd( 'Modulo', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='MODULO'),
                ShdStd( 'Warp', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='WRAP'),
                ShdStd( 'Snap', 'TRACKING_CLEAR_FORWARDS', 'VectorMath', operation='SNAP'),
            ],
            'Mix': [
                ShdStd( 'Mix', 'IMAGE', 'MixRGB', blend_type='MIX' ),
                ShdStd( 'Darken', 'IMAGE', 'MixRGB', blend_type='DARKEN'),
                ShdStd( 'Multiply', 'IMAGE', 'MixRGB', blend_type='MULTIPLY'),
                ShdStd( 'Burn', 'IMAGE', 'MixRGB', blend_type='BURN'),
                (),
                ShdStd( 'Lighten', 'IMAGE', 'MixRGB', blend_type='LIGHTEN'),
                ShdStd( 'Screen', 'IMAGE', 'MixRGB', blend_type='SCREEN'),
                ShdStd( 'Dodge', 'IMAGE', 'MixRGB', blend_type='DODGE'),
                ShdStd( 'Add', 'IMAGE', 'MixRGB', blend_type='ADD'),
                (),
                ShdStd( 'Overlay', 'IMAGE', 'MixRGB', blend_type='OVERLAY'),
                ShdStd( 'Soft Light', 'IMAGE', 'MixRGB', blend_type='SOFT_LIGHT'),
                ShdStd( 'Linear Light', 'IMAGE', 'MixRGB', blend_type='LINEAR_LIGHT'),
                (),
                ShdStd( 'Difference', 'IMAGE', 'MixRGB', blend_type='DIFFERENCE'),
                ShdStd( 'Subract', 'IMAGE', 'MixRGB', blend_type='SUBTRACT'),
                ShdStd( 'Divide', 'IMAGE', 'MixRGB', blend_type='DIVIDE'),
                (),
                ShdStd( 'Hue', 'IMAGE', 'MixRGB', blend_type='HUE'),
                ShdStd( 'Saturation', 'IMAGE', 'MixRGB', blend_type='SATURATION'),
                ShdStd( 'Color', 'IMAGE', 'MixRGB', blend_type='COLOR'),
                ShdStd( 'Value', 'IMAGE', 'MixRGB', blend_type='VALUE'),
            ]
        }

        block_colors = {
            'Utils': colors[0],
            'Info': colors[1],
            'Texture': colors[2],
            'Color': colors[3],
            'Normal': colors[5],
            'Shader': colors[6],
            'Input': colors[4],
            'Math': colors[14],
            'VMath': colors[15],
            'Mix': colors[7],
        }

        columns = [
            [ 'Info', 'Texture', 'Color', 'Shader', 'Input', 'Normal', 'Utils', ],
            [ 'Math' ],
            [ 'VMath' ],
            [ 'Mix' ],
        ]

        l.ui_units_x = 25
        l.ui_units_y = 20
        c = l.grid_flow(columns=len(columns), even_columns=False)

        for _, col in enumerate(columns):
            r = c.column(align=True)
            prev = None
            for k in col:
                #r.label(text=f'{i}{k}:')
                for node in nodes[k]:
                    if isinstance(node, ShdStd):
                        self.__std_simple_shader(r, node, color=block_colors[k])
                    else:
                        # Prevent double seperators (due to do_show flag).
                        if isinstance(prev, ShdStd):
                            r.separator(factor=0.5)
                    prev = node
                r.separator()


    def __new_image(self, l: UILayout, text: str, resolution: int, mode: str, color):
        props = Properties.get()
        op = l.operator(UI_OT_create_node_new_image.bl_idname, text=text)
        op.resolution = resolution
        op.mode = mode
        op.uv_map = props.create_image_uv_map
        op.use_alpha = props.create_image_alpha
        op.use_32bit = props.create_image_32bit
        op.fill = color


    def __draw_pie_textures(self, layout: UILayout):
        props = Properties.get()
        obj = bpy.context.active_object

        c = layout
        s = c.column(align=True)

        if local_addon_info.vl():
            s.label(text='Texture Image')
            if TextureRegistry.instance().get_category_enums(RepositoryType.PBR):
                s.operator('wm.call_menu_pie', text='Create PBR Setup', icon='NODE_MATERIAL').name = NODE_MT_snw_pbr_category_pie_menu.bl_idname
            else:
                s.box().label(text='No PBR categories found')
            if TextureRegistry.instance().get_category_enums(RepositoryType.Image):
                s.operator('wm.call_menu_pie', text='Create Image', icon='IMAGE').name = NODE_MT_snw_image_category_pie_menu.bl_idname
            else:
                s.box().label(text='No Image categories found')
            if TextureRegistry.instance().get_category_enums(RepositoryType.Gray):
                s.operator('wm.call_menu_pie', text='Create Grunge', icon='NODE_COMPOSITING').name = NODE_MT_snw_grunge_category_pie_menu.bl_idname
            else:
                s.box().label(text='No Grunge categories found')
        else:
            b = s.box()
            b.label(text='Trial time expired, texture features disabled.')
            b.label(text='See side panel for info.')


        s.separator()
        s.label(text='New Image')
        if obj.data.uv_layers:
            s.prop(props, 'create_image_uv_map', text='')
            s.prop(props, 'create_image_alpha', toggle=True, icon='IMAGE_ALPHA')
            s.prop(props, 'create_image_32bit', toggle=True, icon='LIGHTPROBE_GRID' if is_ls_410() else 'LIGHTPROBE_VOLUME')

            r = s.column_flow(columns=6, align=True)
            r.label(text='IM')
            self.__new_image(r, '1K', 1024, 'IM', props.create_image_im_color)
            self.__new_image(r, '2K', 2048, 'IM', props.create_image_im_color)
            self.__new_image(r, '4K', 4096, 'IM', props.create_image_im_color)
            self.__new_image(r, '8K', 8192, 'IM', props.create_image_im_color)
            r.prop(props, 'create_image_im_color', text='')

            r = s.column_flow(columns=6, align=True)
            r.label(text='FX')
            self.__new_image(r, '1K', 1024, 'FX', props.create_image_fx_color)
            self.__new_image(r, '2K', 2048, 'FX', props.create_image_fx_color)
            self.__new_image(r, '4K', 4096, 'FX', props.create_image_fx_color)
            self.__new_image(r, '8K', 8192, 'FX', props.create_image_fx_color)
            r.prop(props, 'create_image_fx_color', text='')
        else:
            s.label(text='Object has no UV Layer')

        if obj and obj.type == 'MESH':
            if hasattr(obj.data, 'vertex_colors'):
                s.separator()
                s.label(text='Vertex Colors')
                r = s.column(align=True)

                r.operator(UI_OT_create_node_vertex_colors.bl_idname, text='New', icon='ADD').cols = ''
                for cols in obj.data.vertex_colors:
                    sp = r.split(factor=0.9, align=True)
                    sp.operator(UI_OT_create_node_vertex_colors.bl_idname, text=cols.name).cols = cols.name

                    op = sp.operator(UI_OT_delete_element.bl_idname, text='', icon='REMOVE')
                    op.uv_map = ''
                    op.vertex_colors = cols.name

            if hasattr(obj.data, 'uv_layers'):
                s.separator()
                s.label(text='UV Layers')
                r = s.column(align=True)

                r.operator(UI_OT_create_node_uv_map.bl_idname, text='New', icon='ADD').uv_map = ''
                for mp in obj.data.uv_layers:
                    sp = r.split(factor=0.9, align=True)

                    sp.operator(UI_OT_create_node_uv_map.bl_idname, text=mp.name).uv_map = mp.name

                    op = sp.operator(UI_OT_delete_element.bl_idname, text='', icon='REMOVE')
                    op.uv_map = mp.name
                    op.vertex_colors = ''


    def __draw_others_menu(self, l: UILayout):
        l.menu('NODE_MT_add', text='All Nodes', icon='NODE') 
        if asset_wizard_pro_enabled():
            op = l.operator('wm.call_menu_pie', text='Asset Wizard Pro', icon_value=IconsRegistry.get_icon('awp-logo-mini'))
            op.name = 'NODE_MT_awp_nodes_pie_menu'


    def draw(self, context):
        pie = self.layout.menu_pie()

        self.__draw_std_nodes_tools(pie.row().column(align=True))
        self.__draw_tools(pie.row().column(align=True))
        self.__draw_others_menu(pie.row().column(align=True))
        self.__draw_pie_textures(pie.row().column(align=True))

