# coding: utf-8
# Copyright (C) 2023 Ilia Kurochkin <brothermechanic@yandex.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 shutil
import time
import bpy
import BakeWrangler
from BakeWrangler.nodes.node_tree import BW_TREE_VERSION

logger = logging.getLogger(__name__)

bw_config_1 = {
    'passes': {
        'diffuse': {
            'bake_cat': 'PBR',
            'bake_subcat': 'ALBEDO',
            'suffix': 'D',
        },
        'nornal': {
            'bake_cat': 'CORE',
            'bake_subcat': 'NORMAL',
            'suffix': 'N',
        },
        'ao': {
            'bake_cat': 'CORE',
            'bake_subcat': 'AO',
            'bake_samples': 16,
            'suffix': 'AO',
        },
        'roughness': {
            'bake_cat': 'CORE',
            'bake_subcat': 'ROUGHNESS',
            'suffix': 'R',
        },
        'metallic': {
            'bake_cat': 'PBR',
            'bake_subcat': 'METALLIC',
            'suffix': 'M',
        },
        'uv': {
            'bake_cat': 'WRANG',
            'bake_subcat': 'ISLANDID',
            'suffix': 'UV',
        },
    },
    'margin': 8,
    'ray_dist': 0.001,
    'bake_samples': 4,
}

bw_config_10 = {
    'passes': {
        'nornal': {
            'bake_cat': 'CORE',
            'bake_subcat': 'NORMAL',
            'suffix': 'N',
        },
        'ao': {
            'bake_cat': 'CORE',
            'bake_subcat': 'AO',
            'bake_samples': 16,
            'suffix': 'AO',
        },
    },
    'margin': 8,
    'ray_dist': 0.01,
    'bake_samples': 4,
}

configs = {'bake_1': bw_config_1, 'bake_10': bw_config_10}


def bw_submit(lowpoly_obj_names,
              bw_config,
              resolution,
              bake_path=None,
              area=None,
              **cg_config):
    ''' Submit session and bake textures with BakeWrangler addon. '''

    if not bake_path:
        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').replace('\\', '/')
        if os.path.exists(bake_path):
            shutil.rmtree(bake_path)
        os.makedirs(bake_path, exist_ok=True)

    # create node tree
    baking_tree_name = ''.join([item[0] for item in bw_config['passes'].keys()])
    tree = bpy.data.node_groups.new(name=baking_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
    tree.use_fake_user = 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'] = bw_config['margin']
    tree.nodes['Mesh Settings']['ray_dist'] = bw_config['ray_dist']
    tree.nodes['Sample Settings'].pinned = True
    tree.nodes['Sample Settings']['bake_samples'] = bw_config['bake_samples']
    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.name = 'Batch'

    node_y_pos = 0
    node_pos = 500
    pass_socket = 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] + [cg_config['midpoly']])
        ]
        img_name = '_'.join(lp.name.split('_')[:-1]) + '_'

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

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

        def bw_pass(node_pos, bake_pass, pass_socket):
            node_pass = tree.nodes.new('BakeWrangler_Bake_Pass')
            node_pass.location = (-200, node_y_pos - node_pos)
            node_img = tree.nodes.new('BakeWrangler_Output_Image_Path')
            node_img.location = (200, node_y_pos - node_pos)

            tree.links.new(node_pass.outputs['Color'], node_img.inputs['Color'])

            node_pass.bake_cat = bake_pass['bake_cat']
            if bake_pass.get('bake_cat') == 'CORE':
                node_pass.bake_core = bake_pass['bake_subcat']
            elif bake_pass.get('bake_cat') == 'PBR':
                node_pass.bake_pbr = bake_pass['bake_subcat']
            elif bake_pass.get('bake_cat') == 'WRANG':
                node_pass.bake_wrang = bake_pass['bake_subcat']
            if bake_pass.get('bake_samples'):
                node_pass.bake_samples = bake_pass['bake_samples']
            node_img.inputs['Color'].suffix = bake_pass['suffix']
            node_img.inputs['Split Output'].disp_path = bake_path
            node_img.inputs['Split Output'].img_name = img_name

            tree.links.new(node_inputs.outputs['Mesh'], node_pass.inputs[1])
            tree.links.new(node_img.outputs['Bake'], node_batch.inputs[pass_socket])

        # bake passes
        for bake_pass in bw_config['passes'].keys():
            bw_pass(node_pos, bw_config['passes'][bake_pass], pass_socket)
            pass_socket += 1
            node_pos += 500

        reroute1 = tree.nodes.new('NodeReroute')
        reroute1.location = (-1000, 0)
        reroute2 = tree.nodes.new('NodeReroute')
        reroute2.location = (2000, 0)
        tree.links.new(reroute1.outputs[0], reroute2.inputs[0])
        node_pos = 500
        node_y_pos -= 4000

    return tree, pass_socket


def bw_bake(lowpoly_obj_names, textures_path, resolution, **cg_config):
    BakeWrangler.register()
    output_paths = []
    for config in configs:
        # 1 pass
        bake_path = os.path.join(textures_path, config).replace('\\', '/')
        if os.path.exists(bake_path):
            shutil.rmtree(bake_path)
        os.makedirs(bake_path, exist_ok=True)
        tree_obj, count = bw_submit(
            lowpoly_obj_names, configs[config], resolution, bake_path, **cg_config)

        node_batch = tree_obj.nodes['Batch']

        bpy.ops.bake_wrangler.bake_pass(tree=tree_obj.name, node=node_batch.name, sock=-1)

        logger.info('Baking pipeline is started!')
        # delay until render will be complete
        # ______________________________
        exit_delay = 180
        start_time = time.time()
        while bpy.context.window_manager.bw_status == 1:
            time.sleep(2)
            if not os.listdir(bake_path):
                if time.time() > start_time + exit_delay:
                    break
            if not len(os.listdir(bake_path)) < count:
                break
        # ______________________________
        assert len(os.listdir(bake_path)) == count, (
            'Warning! Baked only {} of {} textures!'.format(
                len(os.listdir(bake_path)), count))
        logger.info('Baking %s textures is finished!', len(os.listdir(bake_path)))

        output_paths.append(bake_path)

    return output_paths