# 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.
Basic image processing for asset pipeline.
'''
__version__ = '0.1'

import logging
import os
import bpy

logger = logging.getLogger(__name__)


def compose_baked_textures(textures_path, bake_paths, resolution):
    ''' Mix, edit, and inpaint baked textures to production condition. '''
    logger.info('Composing textures started!')
    # cleanup
    for img in os.listdir(textures_path):
        if img.startswith('T_'):
            os.remove(os.path.join(textures_path, img).replace('\\', '/'))

    bake_path_1, bake_path_10, = bake_paths

    baked_images = sorted(os.listdir(bake_path_1))

    bpy.context.scene.render.resolution_x = resolution
    bpy.context.scene.render.resolution_y = resolution
    bpy.context.scene.render.resolution_percentage = 100
    bpy.context.scene.frame_start = 1
    bpy.context.scene.frame_end = 1
    bpy.context.scene.render.filepath = textures_path
    bpy.context.scene.render.image_settings.file_format = 'PNG'
    bpy.context.scene.render.image_settings.color_mode = 'RGB'
    bpy.context.scene.render.image_settings.compression = 50

    bpy.context.scene.display_settings.display_device = 'sRGB'
    bpy.context.scene.view_settings.view_transform = 'Standard'
    bpy.context.scene.view_settings.look = 'None'
    bpy.context.scene.view_settings.exposure = 0
    bpy.context.scene.view_settings.gamma = 1
    bpy.context.scene.view_settings.use_curve_mapping = False

    bpy.context.scene.cycles.use_auto_tile = False

    bpy.context.scene.render.use_compositing = True
    bpy.context.scene.render.dither_intensity = 2
    bpy.context.scene.use_nodes = True

    tree = bpy.context.scene.node_tree

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

    node_y_pos = 0
    # render for eatch uv packs
    uv_images = [image for image in baked_images if 'UV' in image]
    for uv_image in uv_images:
        uv_obj = bpy.data.images.load(os.path.join(bake_path_1, uv_image))
        uv_node = tree.nodes.new(type='CompositorNodeImage')
        uv_node.image = bpy.data.images[uv_obj.name]
        uv_node.location = 0, (node_y_pos + 500)
        key_node = tree.nodes.new(type="CompositorNodeLumaMatte")
        key_node.location = 250, (node_y_pos + 500)
        key_node.limit_max = 0.01
        tree.links.new(uv_node.outputs['Image'], key_node.inputs['Image'])
        de_node = tree.nodes.new(type="CompositorNodeDilateErode")
        de_node.location = 500, (node_y_pos + 500)
        de_node.distance = -2
        tree.links.new(key_node.outputs['Matte'], de_node.inputs['Mask'])

        slot_name = uv_image.split('_UV')[0]
        logger.info('Prepare %s image slot!', slot_name)
        for image in baked_images:
            if '_UV' in image:
                continue
            if not image.startswith(slot_name):
                continue
            if '_AO' in image:
                image_obj = bpy.data.images.load(os.path.join(bake_path_1, image))
                image_node = tree.nodes.new(type='CompositorNodeImage')
                image_node.location = -500, (node_y_pos + 200)
                image_node.image = bpy.data.images[image_obj.name]

                image_obj_double = bpy.data.images.load(os.path.join(bake_path_10, image))
                image_node_double = tree.nodes.new(type='CompositorNodeImage')
                image_node_double.location = -500, (node_y_pos - 200)
                image_node_double.image = bpy.data.images[image_obj_double.name]

                lighten_node = tree.nodes.new(type="CompositorNodeMixRGB")
                lighten_node.location = 0, node_y_pos
                lighten_node.blend_type = 'LIGHTEN'

                tree.links.new(image_node.outputs['Image'], lighten_node.inputs[1])
                tree.links.new(image_node_double.outputs['Image'], lighten_node.inputs[2])

                aa_node = tree.nodes.new(type="CompositorNodeAntiAliasing")
                aa_node.location = 500, node_y_pos
                aa_node.threshold = 0
                tree.links.new(lighten_node.outputs['Image'], aa_node.inputs['Image'])

            elif '_N' in image:
                image_obj = bpy.data.images.load(os.path.join(bake_path_1, image))
                image_node = tree.nodes.new(type='CompositorNodeImage')
                image_node.location = -1000, (node_y_pos + 200)
                image_node.image = bpy.data.images[image_obj.name]

                image_obj_double = bpy.data.images.load(os.path.join(bake_path_10, image))
                image_node_double = tree.nodes.new(type='CompositorNodeImage')
                image_node_double.location = -1000, (node_y_pos - 200)
                image_node_double.image = bpy.data.images[image_obj_double.name]

                channel_node = tree.nodes.new(type="CompositorNodeSeparateColor")
                channel_node.location = -500, (node_y_pos + 200)

                tree.links.new(image_node.outputs['Image'], channel_node.inputs['Image'])

                mix_node = tree.nodes.new(type="CompositorNodeMixRGB")
                mix_node.location = 0, node_y_pos
                mix_node.blend_type = 'MIX'

                tree.links.new(channel_node.outputs['Blue'], mix_node.inputs[0])
                tree.links.new(image_node_double.outputs['Image'], mix_node.inputs[1])
                tree.links.new(image_node.outputs['Image'], mix_node.inputs[2])

                aa_node = tree.nodes.new(type="CompositorNodeAntiAliasing")
                aa_node.location = 500, node_y_pos
                aa_node.threshold = 0
                tree.links.new(mix_node.outputs['Image'], aa_node.inputs['Image'])

            elif '_R' in image:
                image_obj = bpy.data.images.load(os.path.join(bake_path_1, image))
                image_node = tree.nodes.new(type='CompositorNodeImage')
                image_node.location = -600, (node_y_pos + 100)
                image_node.image = bpy.data.images[image_obj.name]

                greater_node = tree.nodes.new(type="CompositorNodeMath")
                greater_node.location = -300, (node_y_pos + 200)
                greater_node.operation = 'GREATER_THAN'
                greater_node.inputs[1].default_value = 0.01

                tree.links.new(image_node.outputs['Image'], greater_node.inputs['Value'])

                mix_node = tree.nodes.new(type="CompositorNodeMixRGB")
                mix_node.location = 0, node_y_pos
                mix_node.blend_type = 'MIX'
                mix_node.inputs[1].default_value = (1, 1, 1, 1)

                tree.links.new(greater_node.outputs['Value'], mix_node.inputs[0])
                tree.links.new(image_node.outputs['Image'], mix_node.inputs[2])

                aa_node = tree.nodes.new(type="CompositorNodeAntiAliasing")
                aa_node.location = 500, node_y_pos
                aa_node.threshold = 0
                tree.links.new(mix_node.outputs['Image'], aa_node.inputs['Image'])

            else:
                image_obj = bpy.data.images.load(os.path.join(bake_path_1, image))
                image_node = tree.nodes.new(type='CompositorNodeImage')
                image_node.location = 0, node_y_pos
                image_node.image = bpy.data.images[image_obj.name]

                aa_node = tree.nodes.new(type="CompositorNodeAntiAliasing")
                aa_node.location = 500, node_y_pos
                aa_node.threshold = 0
                tree.links.new(image_node.outputs['Image'], aa_node.inputs['Image'])

            sa_node = tree.nodes.new(type="CompositorNodeSetAlpha")
            sa_node.location = 1000, node_y_pos
            sa_node.mode = 'REPLACE_ALPHA'
            tree.links.new(aa_node.outputs['Image'], sa_node.inputs['Image'])
            tree.links.new(de_node.outputs['Mask'], sa_node.inputs['Alpha'])

            sa_node = tree.nodes.new(type="CompositorNodeSetAlpha")
            sa_node.location = 1000, node_y_pos
            sa_node.mode = 'REPLACE_ALPHA'
            tree.links.new(aa_node.outputs['Image'], sa_node.inputs['Image'])
            tree.links.new(de_node.outputs['Mask'], sa_node.inputs['Alpha'])

            ip_node = tree.nodes.new(type="CompositorNodeInpaint")
            ip_node.location = 1500, node_y_pos
            ip_node.distance = 512
            tree.links.new(sa_node.outputs['Image'], ip_node.inputs['Image'])

            sc_node = tree.nodes.new(type="CompositorNodeScale")
            sc_node.location = 2000, node_y_pos
            sc_node.space = 'RENDER_SIZE'
            tree.links.new(ip_node.outputs['Image'], sc_node.inputs['Image'])

            if '_D' in image or '_N' in image:
                out_node = tree.nodes.new(type="CompositorNodeOutputFile")
                out_node.location = 3000, node_y_pos
                out_node.file_slots[0].path = 'T_{}{}'.format(
                    slot_name,
                    image.split(slot_name)[1].split('.')[0].lower())
                tree.links.new(sc_node.outputs['Image'], out_node.inputs['Image'])

            else:
                gamma_node = tree.nodes.new(type="CompositorNodeGamma")
                gamma_node.location = 2500, node_y_pos
                gamma_node.inputs[1].default_value = 2.2

                tree.links.new(sc_node.outputs['Image'], gamma_node.inputs['Image'])

                out_node = tree.nodes.new(type="CompositorNodeOutputFile")
                out_node.location = 3000, node_y_pos
                out_node.file_slots[0].path = 'T_{}{}'.format(
                    slot_name,
                    image.split(slot_name)[1].split('.')[0].lower())
                tree.links.new(gamma_node.outputs['Image'], out_node.inputs['Image'])

            node_y_pos -= 500

        # render
        bpy.ops.render.render(animation=False)
        # remove all nodes after rendering
        for node in tree.nodes:
            tree.nodes.remove(node)

    # cleanup filenames
    composed_textures = [
        texture
        for texture in os.listdir(textures_path)
        if '000' in texture and texture.endswith('.png')
    ]
    for texture in composed_textures:
        os.rename(
            os.path.join(textures_path, texture).replace('\\', '/'),
            os.path.join(textures_path, texture.rpartition('0001')[0] + '.png').replace('\\', '/')
        )

    logger.info('Composing %s textures is finished!', len(composed_textures))
    return textures_path
