# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility function to convert spec csv to json schema"""
import csv
import json

__type_mapping = {
    'collection': 'object',
    'list': 'array',
    'float': 'number',
    'bool': 'boolean',
    'integer': 'integer',
    'string': 'string',
    'str': 'string',
    'int': 'integer',
    'dict': 'object',
    'const': 'const',
    'ordered': 'ordered',
    'categorical': 'categorical',
    'ordered_int': 'ordered_int',
    'enum': 'string'
}


def __basic_type_fix(value_type, value):
    """Converts spec values based on their datatype"""
    if value in (None, ''):
        return None
    if value in ('inf', '-inf'):
        return float(value)
    if value_type in ('integer', 'ordered_int'):
        return int(value)
    if value_type == 'number':
        return float(value)
    if value_type == 'boolean':
        return str(value).lower() == "true"
    if value_type == 'array':
        return json.loads(value)
    if value_type == 'object':
        return json.loads(value)
    return value


def __array_type_fix(value_type, value):
    """Converts spec values in an array based on their datatype"""
    if value in (None, ''):
        return None
    # We dont need this for list / dict (examples, valid_options only for simple types)
    # if value_type in ["array","object"]:
    #     return None
    values = value.replace(' ', '').split(',')
    if value_type == 'integer':
        return [int(i) for i in values]
    if value_type == 'number':
        return [float(i) for i in values]
    if value_type == 'boolean':
        return [str(i).lower() == "true" for i in values]
    if value_type == 'array':
        return [json.loads(i) for i in values]
    if value_type == 'object':
        return [json.loads(i) for i in values]
    return values


def __merge(d1, d2):
    """Merges two dictionaries"""
    for key in d2.keys():
        if key not in d1:
            d1[key] = d2[key]
        elif d1[key] is None:
            d1[key] = d2[key]
        elif type(d1[key]) is list and type(d2[key]) is list:
            if d1[key] != [] and type(d1[key][0]) is dict:
                for i in range(0, min(len(d1[key]), len(d2[key]))):
                    __merge(d1[key][i], d2[key][i])
            else:
                d1[key] = d1[key] + [i for i in d2[key] if i not in d1[key]]
        elif type(d2[key]) is not dict:
            d1[key] = d2[key]
        else:
            __merge(d1[key], d2[key])
    return d1


def harden_parameter_name(parameter):
    """Fix parameter names by removing flanking "." and remove all spaces"""
    if not parameter:
        return None
    return parameter.rstrip(".").lstrip(".").replace(" ", "")


def harden_value_type(value_type):
    """If value type is an unknown, make it string"""
    if value_type not in __type_mapping.keys():
        return "string"
    return value_type


def harden_numerical_value(value):
    """If the value cannot become a float, then return None"""
    if value:
        try:
            float(value)
            return value
        except:
            return None
    else:
        return None


def convert(path, classes=[]):
    """Convert csv spec to json schema"""
    array_parameters = []
    schema = {}

    with open(path, mode='r', encoding='utf-8-sig') as f:
        reader = csv.DictReader(f)
        for row in reader:
            # get row data
            parameter = row.get('parameter')
            # if parameter:
            #     parameter = harden_parameter_name(parameter)
            display_name = row.get('display_name')
            value_type = row.get('value_type')
            # if value_type:
            #     value_type = harden_value_type(value_type)
            description = row.get('description')
            default_value = row.get('default_value')
            examples = row.get('examples')
            valid_min = row.get('valid_min')
            # if valid_min:
            #     valid_min = harden_numerical_value(valid_min)
            valid_max = row.get('valid_max')
            # if valid_max:
            #     valid_max = harden_numerical_value(valid_max)
            valid_options = row.get('valid_options')
            required = row.get('required')
            popular = row.get('popular')
            automl_enabled = row.get('automl_enabled')
            regex = row.get('regex')
            link = row.get('link')
            # convert value type
            value_type = __type_mapping.get(value_type)
            if value_type is None:
                continue
            if value_type == 'array':
                array_parameters.append(parameter)
            # fix data types
            default_value = __basic_type_fix(value_type, default_value)
            valid_min = __basic_type_fix(value_type, valid_min)
            valid_max = __basic_type_fix(value_type, valid_max)
            valid_options = __array_type_fix(value_type, valid_options)
            examples = __array_type_fix(value_type, examples)
            # compose object
            params = parameter.split('.')
            last_param = params.pop()
            if value_type == 'const':
                obj = {'type': 'object',
                       'properties': {last_param: {'const': default_value}}, 'default': {last_param: default_value}}
            else:
                obj = {'type': 'object',
                       'properties': {last_param: {'type': value_type}}}
                # add known object details
                props = obj['properties'][last_param]
                if display_name not in (None, ''):
                    props['title'] = display_name
                if description not in (None, ''):
                    props['description'] = description
                if examples not in (None, []):
                    props['examples'] = examples
                if default_value not in (None, ''):
                    props['default'] = default_value
                    obj['default'] = {last_param: default_value}
                if valid_min is not None:
                    props['minimum'] = valid_min
                if valid_max is not None:
                    props['maximum'] = valid_max
                if valid_options not in (None, []):
                    props['enum'] = valid_options
                if regex not in (None, '') and value_type == 'string':
                    props['pattern'] = regex
                if link is not None and link.startswith('http'):
                    props['link'] = link
                if required is not None and required.lower() == 'yes':
                    if obj.get('required') is None:
                        obj['required'] = []
                    obj['required'].append(last_param)
                if popular is not None and popular.lower() == 'yes':
                    if obj.get('popular') is None:
                        obj['popular'] = []
                    obj['popular'].append(last_param)
                if automl_enabled is not None and automl_enabled.lower() == 'true':
                    if obj.get('automl_default_parameters') is None:
                        obj['automl_default_parameters'] = []
                    obj['automl_default_parameters'].append(parameter)
                # special override of default with array of class strings
                isArray = parameter in array_parameters
                if classes != [] and isArray:
                    if parameter == 'inferencer_config.target_classes':
                        obj['default'] = {last_param: classes}
            # add object hierarchy
            while len(params) > 0:
                joined_params = '.'.join(params)
                isArray = joined_params in array_parameters
                isRequired = obj.get('required') is not None
                isPopular = obj.get('popular') is not None
                hasDefault = obj.get('default') is not None
                isAutomlenabled = obj.get('automl_default_parameters') is not None
                param = params.pop()
                if isArray:
                    default = []
                    if hasDefault:
                        default = [obj['default']]
                    if classes != []:
                        # dynamic patching of default for given dataset classes
                        if joined_params == 'classwise_config':
                            if hasDefault:
                                default = [__merge({'key': c}, obj['default']) for c in classes]
                            else:
                                default = [{'key': c} for c in classes]
                        elif joined_params == 'bbox_handler_config.classwise_bbox_handler_config':
                            if hasDefault:
                                default = [__merge({'key': c, 'value': {'output_map': c}}, obj['default']) for c in classes]
                            else:
                                default = [{'key': c, 'value': {'output_map': c}} for c in classes]
                        elif joined_params == 'dataset_config.target_class_mapping':
                            default = [{'key': c, 'value': c} for c in classes]
                    obj = {
                        'type': 'object',
                        'properties': {
                            param: {
                                'type': 'array',
                                'items': obj,
                                'default': default
                            }
                        }
                    }
                    if hasDefault or default != []:
                        obj['default'] = {param: default}
                else:
                    default = obj.get('default')
                    obj = {
                        'type': 'object',
                        'properties': {param: obj}
                    }
                    if hasDefault:
                        obj['default'] = {param: default}
                if isRequired:
                    obj['required'] = [param]
                if isPopular:
                    obj['popular'] = [param]
                if isAutomlenabled:
                    if not schema.get('automl_default_parameters'):
                        schema['automl_default_parameters'] = [parameter]
                    else:
                        schema['automl_default_parameters'].append(parameter)
            # update schema with obj
            __merge(schema, obj)

    # return json schema
    return schema
