# coding: utf-8
# Copyright (C) 2023 Ilia Kurochkin <brothermechanic@gmail.com>
#
# 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.
'''
DESCRIPTION.
Preparing and execution methods for the baking process.
'''
__version__ = '0.2'

import logging
import os
import bpy
import addon_utils
import BakeWrangler
from BakeWrangler.nodes.node_tree import BW_TREE_VERSION

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# COLLECTIONS NAMIG CONVENTION
parts_col_name = 'Parts'
lcs_col_name = 'LCS'
hierarchy_col_name = 'Hierarchy'
lowpoly_col_name = 'Lowpoly'
# LCS POINT'S SUFFIXES CONVENTION
inlet = '_in'
outlet = '_out'
root = '_root'
# CG ASSETS SUFFIXES CONVENTION
hightpoly = '_hp'
midpoly = 'mp'
lowpoly = '_lp'
render = '_render'


def bw_submit(lowpoly_obj_names, resolution=4096, tree_name='robossembler', area=None):
    ''' Submit session and bake textures with BakeWrangler addon. '''
    BakeWrangler.register()

    asm_name = os.path.basename(bpy.context.blend_data.filepath).split('.')[0]
    asm_path = os.path.dirname(bpy.context.blend_data.filepath)
    textures_path = os.path.join(asm_path,'textures').replace('\\', '/')
    bake_path = os.path.join(textures_path,'bake')
    bake_path_double = os.path.join(textures_path,'bake_double')
    os.makedirs(bake_path, exist_ok=True)

    # create node tree
    tree = bpy.data.node_groups.new(name=tree_name, type='BakeWrangler_Tree')

    # for default used Compositing Node Editor area
    if not area:
        if 'Compositing' in bpy.data.workspaces:
            for target_area in bpy.data.workspaces['Compositing'].screens[0].areas:
                if target_area.type == 'NODE_EDITOR':
                    area = target_area
                    break
        elif 'Rendering' in bpy.data.workspaces:
            # only if fail for Compositing Workspace
            for target_area in bpy.data.workspaces['Rendering'].screens[0].areas:
                if target_area.type == 'IMAGE_EDITOR':
                    area = target_area
                    break
        else:
            logger.info('Please, set default Bledner settings or set <area> parameter!')
    area.spaces[0].tree_type = 'BakeWrangler_Tree'
    area.spaces[0].path.start(tree)

    tree.tree_version = BW_TREE_VERSION
    tree.initialised = True

    # clear default nodes
    for node in tree.nodes:
        tree.nodes.remove(node)

    # bw settings
    tree.nodes.new('BakeWrangler_MeshSettings').location = (0, 0)
    tree.nodes.new('BakeWrangler_SampleSettings').location = (200, 0)
    tree.nodes.new('BakeWrangler_PassSettings').location = (400, 0)
    tree.nodes.new('BakeWrangler_OutputSettings').location = (600, 0)

    tree.nodes['Mesh Settings'].pinned = True
    tree.nodes['Mesh Settings']['margin'] = 8
    tree.nodes['Mesh Settings']['ray_dist'] = 0.001
    tree.nodes['Sample Settings'].pinned = True
    tree.nodes['Sample Settings']['bake_samples'] = 4
    tree.nodes['Pass Settings'].pinned = True
    tree.nodes['Pass Settings']['res_bake_x'] = resolution
    tree.nodes['Pass Settings']['res_bake_y'] = resolution
    tree.nodes['Output Settings'].pinned = True
    tree.nodes['Output Settings']['img_xres'] = resolution
    tree.nodes['Output Settings']['img_yres'] = resolution
    tree.nodes['Output Settings']['img_type'] = 2  # PNG
    tree.nodes['Output Settings']['img_compression'] = 50
    tree.nodes['Output Settings']['img_color_mode'] = 1  # RGB

    # batch bake node
    node_batch = tree.nodes.new('BakeWrangler_Output_Batch_Bake')
    node_batch.location = (1000, -500)

    node_batch_double = tree.nodes.new('BakeWrangler_Output_Batch_Bake')
    node_batch_double.location = (2000, -1500)

    node_y_pos = 0
    pass_socket = 0
    pass_socket_double = 0
    for lp_name in lowpoly_obj_names:
        # run for eatch lowpoly object
        lp = bpy.data.objects[lp_name]
        mp = bpy.data.objects['_'.join(lp.name.split('_')[:-1] + [midpoly])]
        img_name = '_'.join(lp.name.split('_')[:-1]) + '_'

        node_inputs = tree.nodes.new('BakeWrangler_Bake_Mesh')
        node_inputs.location = (-700, node_y_pos - 500)

        node_inputs.inputs['Target'].value = lp
        node_inputs.inputs['Source'].value = mp

        # bake passes
        # Diffuse
        node_d_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_d_pass.location = (-200, node_y_pos - 500)
        node_d_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_d_img.location = (200, node_y_pos - 500)

        tree.links.new(node_d_pass.outputs['Color'], node_d_img.inputs['Color'])

        node_d_pass.bake_cat = 'PBR'
        node_d_pass.bake_pbr = 'ALBEDO'
        node_d_img.inputs['Color'].suffix = 'D'
        node_d_img.inputs['Split Output'].disp_path = bake_path
        node_d_img.inputs['Split Output'].img_name = img_name
        '''
        ## Curvature
        node_c_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_c_pass.location = (-200, node_y_pos -1000)
        node_c_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_c_img.location = (200, node_y_pos -1000)

        tree.links.new(node_c_pass.outputs['Color'], node_c_img.inputs['Color'])

        node_c_pass.bake_cat = 'WRANG'
        node_c_pass.bake_wrang = 'CURVATURE'
        node_c_img.inputs['Color'].suffix = 'C'
        node_c_img.inputs['Split Output'].disp_path = bake_path
        node_c_img.inputs['Split Output'].img_name = img_name
        '''
        # Normal
        node_n_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_n_pass.location = (-200, node_y_pos - 1500)
        node_n_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_n_img.location = (200, node_y_pos - 1500)

        node_n_img_double = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_n_img_double.location = (200, node_y_pos - 1750)

        tree.links.new(node_n_pass.outputs['Color'], node_n_img.inputs['Color'])
        tree.links.new(node_n_pass.outputs['Color'], node_n_img_double.inputs['Color'])

        node_n_pass.bake_cat = 'CORE'
        node_n_pass.bake_core = 'NORMAL'
        node_n_img.inputs['Color'].suffix = 'N'
        node_n_img.inputs['Split Output'].disp_path = bake_path
        node_n_img.inputs['Split Output'].img_name = img_name

        node_n_img_double.inputs['Color'].suffix = 'N'
        node_n_img_double.inputs['Split Output'].disp_path = bake_path_double
        node_n_img_double.inputs['Split Output'].img_name = img_name

        # AO
        node_ao_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_ao_pass.location = (-200, node_y_pos - 2000)
        node_ao_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_ao_img.location = (200, node_y_pos - 2000)

        node_ao_img_double = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_ao_img_double.location = (200, node_y_pos - 2250)

        tree.links.new(node_ao_pass.outputs['Color'], node_ao_img.inputs['Color'])
        tree.links.new(node_ao_pass.outputs['Color'], node_ao_img_double.inputs['Color'])

        node_ao_pass.bake_cat = 'CORE'
        node_ao_pass.bake_core = 'AO'
        node_ao_pass.bake_samples = 32
        node_ao_img.inputs['Color'].suffix = 'AO'
        node_ao_img.inputs['Split Output'].disp_path = bake_path
        node_ao_img.inputs['Split Output'].img_name = img_name

        node_ao_img_double.inputs['Color'].suffix = 'AO'
        node_ao_img_double.inputs['Split Output'].disp_path = bake_path_double
        node_ao_img_double.inputs['Split Output'].img_name = img_name

        # Roughness
        node_r_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_r_pass.location = (-200, node_y_pos - 2500)
        node_r_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_r_img.location = (200, node_y_pos - 2500)

        tree.links.new(node_r_pass.outputs['Color'], node_r_img.inputs['Color'])

        node_r_pass.bake_cat = 'CORE'
        node_r_pass.bake_core = 'ROUGHNESS'
        node_r_img.inputs['Color'].suffix = 'R'
        node_r_img.inputs['Split Output'].disp_path = bake_path
        node_r_img.inputs['Split Output'].img_name = img_name

        # Metallic
        node_m_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_m_pass.location = (-200, node_y_pos - 3000)
        node_m_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_m_img.location = (200, node_y_pos - 3000)

        tree.links.new(node_m_pass.outputs['Color'], node_m_img.inputs['Color'])

        node_m_pass.bake_cat = 'PBR'
        node_m_pass.bake_pbr = 'METALLIC'
        node_m_img.inputs['Color'].suffix = 'M'
        node_m_img.inputs['Split Output'].disp_path = bake_path
        node_m_img.inputs['Split Output'].img_name = img_name

        # UV
        node_uv_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
        node_uv_pass.location = (-200, node_y_pos - 3500)
        node_uv_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
        node_uv_img.location = (200, node_y_pos - 3500)

        tree.links.new(node_uv_pass.outputs['Color'], node_uv_img.inputs['Color'])

        node_uv_pass.bake_cat = 'WRANG'
        node_uv_pass.bake_wrang = 'ISLANDID'
        node_uv_img.inputs['Color'].suffix = 'UV'
        node_uv_img.inputs['Split Output'].disp_path = bake_path
        node_uv_img.inputs['Split Output'].img_name = img_name

        # connect meshes to passes
        tree.links.new(node_inputs.outputs['Mesh'], node_d_pass.inputs[1])
        '''
        tree.links.new(node_inputs.outputs['Mesh'], node_c_pass.inputs[1])
        '''
        tree.links.new(node_inputs.outputs['Mesh'], node_n_pass.inputs[1])
        tree.links.new(node_inputs.outputs['Mesh'], node_ao_pass.inputs[1])
        tree.links.new(node_inputs.outputs['Mesh'], node_r_pass.inputs[1])
        tree.links.new(node_inputs.outputs['Mesh'], node_m_pass.inputs[1])
        tree.links.new(node_inputs.outputs['Mesh'], node_uv_pass.inputs[1])

        # batch bake node
        tree.links.new(node_d_img.outputs['Bake'], node_batch.inputs[pass_socket])
        '''
        pass_socket += 1
        tree.links.new(node_c_img.outputs['Bake'], node_batch.inputs[pass_socket])
        '''
        pass_socket += 1
        tree.links.new(node_n_img.outputs['Bake'], node_batch.inputs[pass_socket])
        pass_socket += 1
        tree.links.new(node_ao_img.outputs['Bake'], node_batch.inputs[pass_socket])
        pass_socket += 1
        tree.links.new(node_r_img.outputs['Bake'], node_batch.inputs[pass_socket])
        pass_socket += 1
        tree.links.new(node_m_img.outputs['Bake'], node_batch.inputs[pass_socket])
        pass_socket += 1
        tree.links.new(node_uv_img.outputs['Bake'], node_batch.inputs[pass_socket])
        pass_socket += 1

        # batch bake node double
        tree.links.new(node_n_img_double.outputs['Bake'],
                       node_batch_double.inputs[pass_socket_double])
        pass_socket_double += 1
        tree.links.new(node_ao_img_double.outputs['Bake'],
                       node_batch_double.inputs[pass_socket_double])
        pass_socket_double += 1

        node_y_pos -= 4000

    bpy.ops.bake_wrangler.bake_pass(tree=tree_name, node=node_batch.name, sock=-1)
    logger.info('Baking first pass is finished!')

    # double pass
    tree.nodes['Mesh Settings']['ray_dist'] = 0.01

    bpy.ops.bake_wrangler.bake_pass(tree=tree_name, node=node_batch_double.name, sock=-1)
    logger.info('Baking double pass is finished!')

    return textures_path