##############################################################################################################################
#                                                                                                                            #
#  Helper Script to generate a base config block for DeepVA                                                                  #
#                                                                                                                            #
#  Parsed from deepva_modules.json (result of https://api.deepva.com/api/v1/modules)                                         #
#  and deepva_models.json (result of https://api.deepva.com/api/v1/models)                                                   #
#                                                                                                                            #
#  Only modules that are useful for this plugin (image recognition) are included                                             #
#  Useful modules are:                                                                                                       #
#    - landmark_recognition                                                                                                  #
#    - object_scene_recognition    -> recommended model: zero-shot                                                           #
#    - face_recognition                                                                                                      #
#    - lower_third_recognition                                                                                               #
#    - advanced_diversity_analysis                                                                                           #
#    - qrcode_detection                                                                                                      #
#    - face_attributes                                                                                                       #
#                                                                                                                            #
#  The output of this script in deepva_baseconfig.yml must be copied into manifest.yml                                       #
#  at the indicated spot under: base_config.#auto_keyworder#.parameters.configs_deepva                                       #
#                                                                                                                            #
#  The output of this script in deepva_l10n.csv must be copied into the google l10n table                                    #
#  (https://docs.google.com/spreadsheets/d/1glXObMmIUd0uXxdFdiPWRZPLCx6qEUaxDfNnmttave4/edit?gid=2140551706#gid=2140551706)  #
#                                                                                                                            #
##############################################################################################################################

import json
import csv


if __name__ == '__main__':

    INDENT_YML = '  ' * 5

    L10N_PREFIX_PARAMETER = 'server.config.parameter.system.auto_keyworder'
    L10N_PREFIX_OPTION = 'server.config.option.system.auto_keyworder'
    L10N_PREFIX_CONFIGS = 'configs_deepva'

    # this must be updated when there are changes in modules in the DeepVA API
    useful_modules = [
        'landmark_recognition',
        'object_scene_recognition',
        'face_recognition',
        'lower_third_recognition',
        'advanced_diversity_analysis',
        'qrcode_detection',
        'face_attributes',
    ]

    fn_deepva_modules = 'deepva_modules.json'
    fn_deepva_models = 'deepva_models.json'
    fn_l10n_csv = 'deepva_l10n.csv'
    fn_baseconfig_yml = 'deepva_baseconfig.yml'

    loaded_modules = None
    with open(fn_deepva_modules, 'r') as f:
        loaded_modules = json.load(fp=f)

    if not isinstance(loaded_modules, list):
        print(f'expected array in {fn_deepva_modules}')
        exit(1)

    loaded_models = None
    with open(fn_deepva_models, 'r') as f:
        loaded_models = json.load(fp=f)

    if not isinstance(loaded_models, dict):
        print(f'expected array in {fn_deepva_modules} "data"')
        exit(1)

    loaded_models = loaded_models.get('data')
    if not isinstance(loaded_models, list):
        print(f'expected array in {fn_deepva_modules} "data"')
        exit(1)

    models_by_type = {}
    for model in loaded_models:
        name = model.get('name')
        if not isinstance(name, str):
            continue
        typ = model.get('type')
        if not isinstance(typ, str):
            continue

        description = model.get('description')
        if not (
            isinstance(description, str) and description != '' and description != '-'
        ):
            description = None

        if not typ in models_by_type:
            models_by_type[typ] = []
        models_by_type[typ].append(
            {
                'name': name,
                'description': description,
            }
        )

    print('models_by_type:')
    print(json.dumps(models_by_type, indent=2))

    yaml_lines = []
    l10n_lines = []

    position = 9  # last position in current baseconfig

    for module in sorted(loaded_modules, key=lambda m: m['name']):
        name = module.get('name')
        if not isinstance(name, str):
            continue

        if name not in useful_modules:
            continue

        position += 1

        # name of module = checkbox if enabled/disabled
        yaml_lines.append(f'- name: module__{name}')
        yaml_lines.append(f'  type: bool')
        yaml_lines.append(f'  default: false')
        yaml_lines.append(f'  position: {position}')

        module_title = name
        title = module.get('title')
        if isinstance(title, str):
            module_title = title
        l10n_lines.append(
            {
                'key': f'{L10N_PREFIX_PARAMETER}.{L10N_PREFIX_CONFIGS}.module__{name}.label',
                'en': f'Module "{module_title}"',
            }
        )

        hint = f"Module `{name}`"

        description = module.get('description')
        if isinstance(description, str) and description != '' and description != '-':
            hint += f': {description}'

        reference_link = module.get('reference_link')
        if (
            isinstance(reference_link, str)
            and reference_link != ''
            and reference_link != '-'
        ):
            hint += f'\n\nSee: {reference_link}'

        l10n_lines.append(
            {
                'key': f'{L10N_PREFIX_PARAMETER}.{L10N_PREFIX_CONFIGS}.module__{name}.label|hint',
                'en': hint,
            }
        )
        l10n_lines.append(
            {
                'key': f'{L10N_PREFIX_PARAMETER}.module__{name}.checkbox',
                'en': 'Enable this module',
            }
        )

        fields = module.get('fields')
        models = []
        recommend_zero_shot = False

        # dropdown with models for this module
        if isinstance(fields, list):
            for field in fields:
                model_type = field.get('model_type')
                if not isinstance(model_type, str):
                    continue

                if model_type not in models_by_type:
                    continue

                module_models = models_by_type.get(model_type)
                if not isinstance(models, list):
                    continue

                for model in module_models:
                    model_name = model.get('name')
                    models.append(model_name)
                    if model_name == 'zero-shot':
                        recommend_zero_shot = True

                    option = f"{model_name}"
                    model_desc = model.get('description')
                    if isinstance(model_desc, str):
                        option += f': {model_desc}'

                    l10n_lines.append(
                        {
                            'key': f'{L10N_PREFIX_OPTION}.model__{name}.{model_name}',
                            'en': option,
                        }
                    )

        # zero-shot model is recommended -> first in list
        models = sorted(models)
        if recommend_zero_shot:
            models.remove('zero-shot')
            models = ['zero-shot'] + models

        if len(models) > 0:
            yaml_lines.append(f'- name: model__{name}')
            yaml_lines.append(f'  type: select')

            l10n_lines.append(
                {
                    'key': f'{L10N_PREFIX_PARAMETER}.{L10N_PREFIX_CONFIGS}.model__{name}.label',
                    'en': f'Model for module "{module_title}"',
                }
            )
            if recommend_zero_shot:
                l10n_lines.append(
                    {
                        'key': f'{L10N_PREFIX_PARAMETER}.{L10N_PREFIX_CONFIGS}.model__{name}.label|hint',
                        'en': 'Model `zero-shot` is recommended',
                    }
                )

            yaml_lines.append(f'  options:')
            for model in models:
                yaml_lines.append(f'  - "{model}"')

            position += 1

            yaml_lines.append(f'  default: "{models[0]}"')
            yaml_lines.append(f'  position: {position}')

    if len(yaml_lines) == 0:
        print(f'expected non-empty array in {fn_deepva_modules}')
        exit(1)

    with open(fn_baseconfig_yml, 'w') as f:
        for row in yaml_lines:
            f.write(INDENT_YML + row + '\n')
        print(f'wrote {len(yaml_lines)} yaml lines into {fn_baseconfig_yml}')

    with open(fn_l10n_csv, 'w') as f:
        writer = csv.DictWriter(
            f,
            ['key', 'tab1', 'tab2', 'en'],
            delimiter='\t',
            quotechar='"',
            escapechar='\\',
            lineterminator='\n',
        )
        for row in l10n_lines:
            writer.writerow(row)
        print(f'wrote {len(l10n_lines)} l10n lines into {fn_l10n_csv}')
