# -*- coding: utf-8 -*-
#!/usr/bin/env python
'''
DESCRIPTION.
CG asset (models) creation pipelines from FreeCAD data.
Support Blender compiled as a Python Module only!
'''
__version__ = '0.8'
import json
import logging
import shutil
import os
from itertools import zip_longest

from blender.utils.cleanup_orphan_data import cleanup_orphan_data
from blender.utils.collection_tools import (copy_collections_recursive,
                                            remove_collections_with_objects)
from utils.cmd_proc import cmd_proc
from blender.import_cad.build_blender_scene import json_to_blend
from blender.processing.restruct_hierarchy import (hierarchy_assembly,
                                                   hierarchy_separated_parts,
                                                   hierarchy_single_part)
from blender.processing.highpoly_setup import setup_meshes
from blender.processing.midpoly_setup import hightpoly_collections_to_midpoly
from blender.processing.lowpoly_setup import parts_to_shells
from blender.processing.uv_setup import uv_unwrap
from blender.texturing.bake_submitter import bw_bake
from blender.texturing.composing import compose_baked_textures
from blender.texturing.shading import assign_pbr_material
from blender.export.dae import export_dae
from blender.export.stl import export_stl
from blender.export.fbx import export_fbx
from blender.export.ply import export_ply
import bpy

# TODO Path
freecad_to_json_script = 'freecad_to_json.py'

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

# TODO NIX
freecad_bin = 'freecadcmd'
# TODO NIX BakwWrangler/nodes/node_tree.py:659
blender_bin = 'blender'

# TODO WEBAPP
cg_config = {
    'lcs_col_name': 'LCS',
    'parts_col_name': 'Parts',
    'midpoly_col_name': 'Midpoly',
    'lowpoly_col_name': 'Lowpoly',
    'lcs_inlet': 'in',
    'lcs_outlet': 'out',
    'lcs_root': 'root',
    'hightpoly': 'hp',
    'midpoly': 'mp',
    'lowpoly': 'lp',
    'render': 'render',
}

defined_pipeline_list = ['cad', 'highpoly', 'midpoly', 'lowpoly', 'baking', 'export']


def cg_pipeline(**kwargs):
    ''' CG assets (models) creation pipeline '''
    # set defaults
    # ______________________________
    assembly_name = (
        kwargs['fcstd_path'].replace('\\', '/').rpartition('/')[2].rpartition('.')[0])
    # clear other paths from kwargs (freecad issue)
    parts_sequence_path = kwargs.pop('parts_sequence_path', None)
    blend_path = kwargs.pop('blend_path', None)
    export_path = kwargs.pop('export_path', None)

    if kwargs['pipeline_type'] not in defined_pipeline_list:
        return logger.error('Pipeline type %s is not defined!', kwargs['pipeline_type'])

    # output file management
    if not blend_path:
        blend_path = kwargs['fcstd_path'].replace('\\', '/').rpartition('/')[0]
    if not export_path:
        export_path = os.path.join(blend_path, 'models').replace('\\', '/')
    os.makedirs(blend_path, exist_ok=True)
    #blend_file = os.path.join(blend_path, f'{assembly_name}.blend').replace('\\', '/')
    remove_collections_with_objects()
    cleanup_orphan_data()

    # 1) сonvert FreeCAD scene to Blender scene
    # ______________________________
    cad_objects = json_to_blend(
        json.loads(
            cmd_proc(freecad_bin,
                    freecad_to_json_script,
                    '--',
                    **kwargs
                    ).split('FreeCAD ')[0]
        ), **cg_config
    )

    # Save original cad setup as blender scene
    if kwargs['pipeline_type'] == 'cad':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        return blend_file

    # 2) prepare highpoly (depend of cad_objects['objs_lcs'] and parts_sequence)
    # ______________________________
    if parts_sequence_path and os.path.isfile(parts_sequence_path):
        with open(parts_sequence_path, 'r', encoding='utf-8') as sequence_file:
            parts_sequence = json.load(sequence_file)
    else:
        parts_sequence = None
    lcs_inlet_objects = [
        inlet for inlet in cad_objects['objs_lcs']
        if inlet.endswith(cg_config['lcs_inlet'])]

    # input cases
    # 1 case
    if parts_sequence and len(lcs_inlet_objects) > 1:
        logger.info('Parts assembling sequence and LCS points found! '
                    'Launch "hierarchy_assembly" restructuring pipeline.')
        part_names = hierarchy_assembly(
            cad_objects['objs_lcs'], parts_sequence, **cg_config)
    # 2 case
    elif parts_sequence and len(lcs_inlet_objects) < 2:
        return logger.error('Assembly do not have enough LCS points!')
    # 3 case
    elif not parts_sequence and lcs_inlet_objects:
        logger.info('Parts assembling sequence not found! '
                    'Launch "hierarchy_separated_parts" restructuring pipeline.')
        part_names = hierarchy_separated_parts(
            cad_objects['objs_lcs'], **cg_config)
    # 4 case
    elif not parts_sequence and not lcs_inlet_objects:
        logger.info('Parts assembling sequence and LCS points not found! '
                    'Launch "hierarchy_single_part" restructuring pipeline.')
        part_names = hierarchy_single_part(**cg_config)

    if not part_names:
        return logger.error('Can not generate parts!')

    # setup highpolys with materials only
    if cad_objects['objs_foreground']:
        setup_meshes(cad_objects['objs_foreground'],
                    sharpness=True, shading=True)
    # setup all highpolys
    else:
        setup_meshes(cad_objects['objs_background'],
                    sharpness=True, shading=True)

    # Save highpoly setup as blender scene
    if kwargs['pipeline_type'] == 'highpoly':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s original hightpoly collections ready!', len(part_names))
        return blend_file

    # 3) prepare midpoly
    # ______________________________
    tmp_col_name = copy_collections_recursive(
        bpy.data.collections[cg_config['parts_col_name']],
        suffix=cg_config['midpoly']
    )
    midpoly_obj_names = hightpoly_collections_to_midpoly(
        tmp_col_name, part_names, **cg_config)

    # Save midpoly setup as blender scene
    if kwargs['pipeline_type'] == 'midpoly':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s midpoly objects ready!', len(midpoly_obj_names))
        return blend_file

    # 4) prepare lowpoly
    # ______________________________
    lowpoly_obj_names = parts_to_shells(part_names, **cg_config)
    uv_unwrap(lowpoly_obj_names)

    # Save lowpoly setup as blender scene
    if kwargs['pipeline_type'] == 'lowpoly':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s lowpoly objects ready!', len(lowpoly_obj_names))
        return blend_file

    # 5) bake textures
    # ______________________________
    if kwargs['textures_resolution'] == 0:
        logger.info('Baking pipeline has been canceled!')
    else:
        textures_path = os.path.join(blend_path, 'textures').replace('\\', '/')
        bake_paths = bw_bake(lowpoly_obj_names,
                            textures_path,
                            kwargs['textures_resolution'],
                            **cg_config)
        # Save baking result
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, 'baking')
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        # 5 prepare textures
        bpy.ops.wm.quit_blender()
        compose_baked_textures(textures_path, bake_paths, kwargs['textures_resolution'])
        for bake_path in bake_paths:
            shutil.rmtree(bake_path)
        bpy.ops.wm.open_mainfile(filepath=blend_file)
        assign_pbr_material(lowpoly_obj_names, textures_path)
        bpy.ops.file.make_paths_relative()

    # Save baking setup as blender scene
    if kwargs['pipeline_type'] == 'baking':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)
        logger.info('%s lowpoly objects baked!', len(lowpoly_obj_names))
        return blend_file

    # 6 export object meshes
    # ______________________________
    # TODO asset manager
    # Save blender scene
    if kwargs['pipeline_type'] == 'export':
        blend_file = os.path.join(
            blend_path,
            '{}_{}.blend'.format(assembly_name, kwargs['pipeline_type'])
        ).replace('\\', '/')
        bpy.ops.wm.save_as_mainfile(filepath=blend_file)

    for part_name in part_names:
        export_fbx(
            obj_name=f'{part_name}_{cg_config["lowpoly"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))
        export_ply(
            obj_name=f'{part_name}_{cg_config["midpoly"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))
        export_dae(
            obj_name=f'{part_name}_{cg_config["midpoly"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))
        export_stl(
            obj_name=f'{part_name}_{cg_config["lowpoly"]}',
            path=os.path.join(export_path, part_name, 'meshes').replace('\\', '/'))

    logger.info('%s parts exported!', len(part_names))
    return blend_file


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
         description='Convert and setup FreeCAD solid objects to 3d assets mesh files.'
    )
    parser.add_argument(
        '--fcstd_path',
        type=str,
        help='Path to source FreeCAD scene',
        required=True
    )
    parser.add_argument(
        '--tesselation_method',
        type=str,
        help='Select tesselation method: Standard or FEM.',
        default='Standard',
        required=False
    )
    parser.add_argument(
        '--linear_deflection',
        type=float,
        help='Max linear distance error',
        default=0.1,
        required=False
    )
    parser.add_argument(
        '--angular_deflection',
        type=float,
        help='Max angular distance error',
        default=20.0,
        required=False
    )
    parser.add_argument(
        '--fem_size',
        type=float,
        help='For FEM method only! Finite element size in mm',
        default=50.0,
        required=False
    )
    parser.add_argument(
        '--skiphidden',
        type=bool,
        help='Skip processing for hidden FreeCAD objects',
        default=True,
        required=False
    )
    parser.add_argument(
        '--property_forse_nonsolid',
        type=str,
        help='FreeCAD property to enable processing for nonsolid objects',
        default='Robossembler_NonSolid',
        required=False
    )
    parser.add_argument(
        '--pipeline_type',
        type=str,
        help='Set pipeline type: "cad", "highpoly", "midpoly", "lowpoly", "baking", "export"',
        default='export',
        required=False
    )
    parser.add_argument(
        '--parts_sequence_path',
        type=str,
        help='Path to parts assembling sequence json file.',
        required=False
    )
    parser.add_argument(
        '--export_path',
        type=str,
        help='Path for export assets. If not, fcstd_path will be used instead.',
        required=False
    )
    parser.add_argument(
        '--blend_path',
        type=str,
        help='Path for blender scene. If not, fcstd_path will be used instead.',
        required=False
    )
    parser.add_argument(
        '--textures_resolution',
        type=int,
        help='Set baking texture resolution. Recomended - 4096 pix ',
        default=512,
        required=False
    )

    args = parser.parse_args()

    cg_kwargs = {key: getattr(args, key) for key in dir(args) if not key.startswith('_')}

    cg_pipeline(**cg_kwargs)