# 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.
Reorganization and restructuring of assembly structure.
'''
__version__ = '0.3'
import logging

import bpy
import mathutils

from blender.utils.object_relations import (parenting,
                                            unparenting)
from blender.utils.object_transforms import round_transforms
from blender.utils.collection_tools import unlink_from_collections

logger = logging.getLogger(__name__)


def hierarchy_assembly(lcs_names, parts_sequence):
    ''' Hierarchy by LCS and Parts Assembling Sequence. '''
    # collect scene hierarchy start info
    main_locators = [obj for obj in bpy.data.objects if not obj.parent]
    lcs_inlet_objects = []
    lcs_outlet_objects = []
    for lcs_name in lcs_names:
        lcs_obj = bpy.data.objects[lcs_name]
        if (lcs_obj.get('Robossembler_SocketFlow') == 'inlet'
            and lcs_obj.get('Robossembler_DefaultOrigin')
            ):
            lcs_inlet_objects.append(lcs_obj)
        else:
            lcs_outlet_objects.append(lcs_obj)

    if not lcs_inlet_objects:
        raise Exception('No LCS Inlet objects found!')

    # get main_locator
    main_locator = None
    for locator in main_locators:
        if set(lcs_inlet_objects + lcs_outlet_objects).issubset(
            locator.children_recursive
        ):
            main_locator = locator
    if not main_locator:
        # TODO need checking
        return logger.error('CAD root locator should be parent of all LCS!')

    # check parts_sequence objects in scene
    for part in parts_sequence:
        if not bpy.data.objects.get(part):
            return logger.error('%s part object not found!', part)

    # create root lcs by parts sequence
    first_part_obj = bpy.data.objects[parts_sequence[0]]
    root_lcs = None
    for lcs_inlet in first_part_obj.children:
        # drop non lcs objs
        if lcs_inlet.name not in lcs_names:
            continue
        # drop non DefaultOrigins
        lcs_obj = bpy.data.objects[lcs_name]
        if not (lcs_obj.get('Robossembler_SocketFlow') == 'inlet'
            and lcs_obj.get('Robossembler_DefaultOrigin')
            ):
            continue
        root_lcs_name = 'root'
        root_lcs = bpy.data.objects.new(root_lcs_name, None)
        root_lcs.empty_display_type = 'ARROWS'
        root_lcs.empty_display_size = 0.15
        root_lcs.show_in_front = True
        root_lcs.location = lcs_inlet.location
        root_lcs.rotation_euler = lcs_inlet.rotation_euler
        root_lcs.parent = lcs_inlet.parent
        bpy.context.scene.collection.objects.link(root_lcs)
        logger.info('Root Inlet LCS object created!')
    unparenting(root_lcs)
    round_transforms(root_lcs)
    parenting(root_lcs, main_locator)

    # retree_by lcs
    for lcs_inlet_obj in lcs_inlet_objects:
        # lcs inlet as main parent
        parent_locator = lcs_inlet_obj.parent
        if not parent_locator:
            raise Exception('LCS %s should have a parent!', lcs_inlet_obj.name)
        unparenting(lcs_inlet_obj)
        round_transforms(lcs_inlet_obj)
        if parent_locator:
            if parent_locator.parent:
                unparenting(parent_locator)
            parenting(lcs_inlet_obj, parent_locator)
        parenting(root_lcs, lcs_inlet_obj)
        # lcs outlet parent to lcs inlet
        for lcs_outlet_obj in lcs_inlet_obj.children_recursive:
            if lcs_outlet_obj.name not in lcs_names:
                continue
            unparenting(lcs_outlet_obj)
            round_transforms(lcs_outlet_obj)
            parenting(lcs_inlet_obj, lcs_outlet_obj)

    # reset transforms for root_lcs
    root_lcs.matrix_world = mathutils.Matrix()

    # lcs collections
    part_names = []
    for lcs_inlet_obj in root_lcs.children:
        # remove unmarked parts
        if lcs_inlet_obj not in lcs_inlet_objects:
            for obj in lcs_inlet_obj.children_recursive:
                bpy.data.objects.remove(obj, do_unlink=True)
            bpy.data.objects.remove(lcs_inlet_obj, do_unlink=True)
            continue
        # collect part names
        part_name = None
        for locator in lcs_inlet_obj.children:
            if locator in lcs_outlet_objects:
                continue
            part_name = locator.name
            part_names.append(part_name)
        # pack parts to collections
        part_col = bpy.data.collections.new(f'{part_name}')
        bpy.data.collections['Render'].children.link(part_col)
        for obj in lcs_inlet_obj.children_recursive:
            unlink_from_collections(obj)
            part_col.objects.link(obj)
        unlink_from_collections(lcs_inlet_obj)
        part_col.objects.link(lcs_inlet_obj)

    # TODO DEPRECATED
    """
    # parts assembling
    for idx, part_name in enumerate(parts_sequence):
        #  TODO clones for clones
        if part_name not in part_names:
            continue
        lcs_inlet_obj = bpy.data.objects[part_name].parent
        constraint = lcs_inlet_obj.constraints.new(type='COPY_TRANSFORMS')
        # drop first_part_obj
        if idx == 0:
            constraint.target = root_lcs
            continue
        # if asm pair exists
        part_before = bpy.data.objects.get(parts_sequence[idx - 1])
        if part_before:
            lcs_outlet_objs = [
                lcs_out
                for lcs_out in part_before.parent.children
                if lcs_out in lcs_outlet_objects]
            if lcs_outlet_objs:
                constraint.target = lcs_outlet_objs[0]
        else:
            constraint.target = root_lcs
            constraint.enabled = False
    # for reset transforms when exporting
    for lcs in lcs_outlet_objects:
        constraint = lcs.constraints.new(type='COPY_TRANSFORMS')
        constraint.target = root_lcs
        constraint.enabled = False
    """

    logger.info('Restructuring assembly pipeline finished!')

    return part_names


def hierarchy_separated_parts(lcs_names):
    ''' Restructuring pipeline as separated parts. '''
    # collect scene hierarchy start info
    lcs_inlet_objects = []
    lcs_outlet_objects = []

    for lcs_name in lcs_names:
        lcs_obj = bpy.data.objects[lcs_name]
        if (lcs_obj.get('Robossembler_SocketFlow') == 'inlet'
            and lcs_obj.get('Robossembler_DefaultOrigin')
            ):
            lcs_inlet_objects.append(lcs_obj)
        else:
            lcs_outlet_objects.append(lcs_obj)

    if not lcs_inlet_objects:
        raise Exception('No LCS Inlet objects found!')

    # retree_by lcs
    part_names = []
    for lcs_inlet_obj in lcs_inlet_objects:
        # lcs inlet as main parent
        parent_locator = lcs_inlet_obj.parent
        if not parent_locator:
            raise Exception('LCS %s should have a parent!', lcs_inlet_obj.name)
        unparenting(lcs_inlet_obj)
        round_transforms(lcs_inlet_obj)
        if parent_locator:
            if parent_locator.parent:
                unparenting(parent_locator)
            parenting(lcs_inlet_obj, parent_locator)
        # lcs outlet parent to lcs inlet
        for lcs_outlet_obj in lcs_inlet_obj.children_recursive:
            if lcs_outlet_obj.name not in lcs_names:
                continue
            unparenting(lcs_outlet_obj)
            round_transforms(lcs_outlet_obj)
            parenting(lcs_inlet_obj, lcs_outlet_obj)

        # reset transforms for inlet_lcs
        lcs_inlet_obj.matrix_world = mathutils.Matrix()

        # pack parts to collections
        part_name = None
        for locator in lcs_inlet_obj.children:
            if locator not in lcs_outlet_objects:
                part_name = locator.name
                part_names.append(part_name)
        part_col = bpy.data.collections.new(f'{part_name}')
        bpy.data.collections['Render'].children.link(part_col)
        for obj in lcs_inlet_obj.children_recursive:
            unlink_from_collections(obj)
            part_col.objects.link(obj)
        unlink_from_collections(lcs_inlet_obj)
        part_col.objects.link(lcs_inlet_obj)
    # remove unmarked objects
    marked_objs = sum(
        [lcs_inlet_obj.children_recursive for lcs_inlet_obj in lcs_inlet_objects],
        [])
    parts_col_objs = bpy.data.collections['Render'].objects
    unmarked_objs = list(set(parts_col_objs) - set(marked_objs))
    if unmarked_objs:
        removed_objs = list(map(bpy.data.objects.remove, unmarked_objs))
        logger.info('%s unmarked objects removed!', len(removed_objs))

    logger.info('Restructuring pipeline as separated parts finished!')

    return part_names


def hierarchy_mono_part():
    ''' Restructuring pipeline as single part. '''
    # collect scene hierarchy start info
    main_locators = [obj for obj in bpy.data.objects if not obj.parent]

    # pack parts to collections
    part_names = []
    for main_locator in main_locators:
        part_name = main_locator.name
        part_names.append(part_name)
        part_col = bpy.data.collections.new(f'{part_name}')
        bpy.data.collections['Render'].children.link(part_col)
        for obj in main_locator.children_recursive:
            unlink_from_collections(obj)
            part_col.objects.link(obj)
        unlink_from_collections(main_locator)
        part_col.objects.link(main_locator)

    logger.info('Restructuring pipeline as single part finished!')

    return part_names