#  ***** BEGIN GPL LICENSE BLOCK *****
#
#  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.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, see <http://www.gnu.org/licenses/>
#  and write to the Free Software Foundation, Inc., 51 Franklin Street,
#  Fifth Floor, Boston, MA  02110-1301, USA..
#
#  The Original Code is Copyright (C) 2023 by Kurochkin Ilia  ###
#  All rights reserved.
#
#  Contact:      brothermechanic@yandex.com    ###
#  Information:  https://gitlab.com/robossembler   ###
#
#  The Original Code is: all of this file.
#
#  ***** END GPL LICENSE BLOCK *****
#
# -*- coding: utf-8 -*-
'''
DESCRIPTION.
Collect all root objects in blender scene and export as json scene configuration.
'''
__version__ = '0.1'

import collections
import json
import logging
import math
import os

from .model_paths import get_urdf_sdf_model_paths
from .model_name import get_sdf_urdf_model_name
from .model_md5 import create_file_md5


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


def export_json(context, scene_config_path, physics_engine):
    ''' Export scene config to json. '''
    # open conf file
    if os.path.isfile(scene_config_path):
        with open(scene_config_path, 'r', encoding='utf-8') as conf_file:
            world_dic = json.load(conf_file)
    else:
        world_dic = collections.defaultdict(list)
        models_dir = os.path.join(os.path.dirname(scene_config_path), 'models')
        model_paths = get_urdf_sdf_model_paths(models_dir)
        for model_path in model_paths:
            model_dic = {}
            # 1
            model_dic['name'] = get_sdf_urdf_model_name(model_path)
            # 2
            model_dic['id'] = create_file_md5(model_path)
            # 3
            relative_model_path = model_path.split(
                os.path.dirname(scene_config_path))[1][1:]
            model_dic['path'] = relative_model_path
            world_dic['models'].append(model_dic)

    # collect all non parented objs
    roots = [obj for obj in context.scene.objects if not obj.parent]

    # collect instances
    for instance in roots:
        instance_dic = {}
        # 1
        instance_dic['model_name'] = ''
        if instance.get('model/name'):
            instance_dic['model_name'] = instance['model/name']
        # 2
        instance_dic['model_id'] = ''
        for item in world_dic['models']:
            if item['name'] != instance_dic['model_name']:
                continue
            instance_dic['model_id'] = item['id']
        # 2
        instance_dic['id'] = ''
        # 3
        instance_dic['pose'] = {}
        instance_dic['pose']['x'] = round(instance.location[0], 3)
        instance_dic['pose']['y'] = round(instance.location[1], 3)
        instance_dic['pose']['z'] = round(instance.location[2], 3)
        instance_dic['pose']['roll'] = round(math.degrees(instance.rotation_euler[0]), 3)
        instance_dic['pose']['pitch'] = round(math.degrees(instance.rotation_euler[1]), 3)
        instance_dic['pose']['yaw'] = round(math.degrees(instance.rotation_euler[2]), 3)
        # 4
        if instance.scale[0] != sum(instance.scale[:]) / 3:
            logger.warning('Asset instance {} is not uniformly scaled!'.format(
                instance.name))
        instance_dic['scale'] = round(instance.scale[0], 3)
        # 5
        # if assembly
        if instance_dic['model_name']:
            instance_dic['type'] = 'asset'
            links = instance.children_recursive
            links.insert(0, instance)
            instance_dic['link_names'] = []
            instance_dic['links'] = {}
            transforms = 0.0
            for link in links:
                child = {}
                if link.type != 'ARMATURE':
                    continue
                bone = link.pose.bones[0]
                rotation_mode = bone.rotation_mode
                bone.rotation_mode = 'XYZ'
                child['pose'] = {}
                child['pose']['x'] = round(bone.location[0], 6)
                child['pose']['y'] = round(bone.location[1], 6)
                child['pose']['z'] = round(bone.location[2], 6)
                child['pose']['roll'] = round(math.degrees(bone.rotation_euler[0]), 6)
                child['pose']['pitch'] = round(math.degrees(bone.rotation_euler[1]), 6)
                child['pose']['yaw'] = round(math.degrees(bone.rotation_euler[2]), 6)
                bone.rotation_mode = rotation_mode
                if bone.scale[0] != sum(bone.scale[:]) / 3:
                    logger.warning(
                        'Link {} of asset {} is not uniformly scaled!'.format(
                            link.name, asset_dic['name']))
                child['scale'] = round(bone.scale[0], 3)
                transforms += round(sum(
                    [child['pose']['x'], child['pose']['y'], child['pose']['z'],
                     child['pose']['roll'], child['pose']['pitch'], child['pose']['yaw']],
                    (child['scale'] - 1.0)),
                    6)
                if '.' in link.name:
                    link_name = link.name.rpartition('.')[0]
                else:
                    link_name = link.name

                instance_dic['link_names'].append(link_name)
                instance_dic['links'][link_name] = child

            if transforms == 0.0:
                instance_dic.pop('link_names')
                instance_dic.pop('links')

            # 6
            instance_dic['parent'] = 'world'
        # if light
        if instance.type == 'LIGHT':
            instance_dic['type'] = instance.type.lower()
            if instance.data.type == 'POINT':
                instance_dic['light_type'] = 'point'
            elif instance.data.type == 'SUN':
                instance_dic['light_type'] = 'directional'
            elif instance.data.type == 'SPOT':
                instance_dic['light_type'] = 'spot'
                instance_dic['spot_angle'] = round(
                    math.degrees(instance.data.spot_size), 1)
            else:
                logger.warning(
                    'Unsupported light type {} on instance {}!'.format(
                        instance.data.type, instance.name))
            instance_dic['intencity'] = instance.data.energy
            instance_dic['diffuse'] = list(instance.data.color)
        # if camera
        if instance.type == 'CAMERA':
            instance_dic['type'] = instance.type.lower()
            instance_dic['focal_length'] = instance.data.lens
            instance_dic['sensor_size'] = [instance.data.sensor_width,
                                           instance.data.sensor_height]

        world_dic['instances'].append(instance_dic)

        # select physics engine
        world_dic['physics'] = {}
        world_dic['physics']['engine_name'] = physics_engine
        world_dic['physics']['gravity'] = {}
        world_dic['physics']['gravity']['x'] = round(context.scene.gravity[0], 3)
        world_dic['physics']['gravity']['y'] = round(context.scene.gravity[1], 3)
        world_dic['physics']['gravity']['z'] = round(context.scene.gravity[2], 3)

    # write conf file
    with open(scene_config_path, 'w', encoding='utf-8') as world_conf_file:
        json.dump(world_dic, world_conf_file, ensure_ascii=False, indent=4)
    logger.info('Scene configuration in json file was completed successfully!')
    return scene_config_path
