#!/usr/bin/env bash
''''export SCRIPT_PATH="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)" # '''
''''export PY_PATH=${SCRIPT_PATH}/venv/bin/python                                       # '''
''''test -f ${PY_PATH} && exec ${PY_PATH} "$0" "$@"                                     # '''
''''echo "Error: Virtual environment not found. Please run 'one-prepare-venv' command." # '''
''''exit 255                                                                            # '''

# Copyright (c) 2021 Samsung Electronics Co., Ltd. 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.

import argparse
import configparser
import os
import subprocess
import sys

from onelib.CfgRunner import CfgRunner
from onelib.WorkflowRunner import WorkflowRunner
import onelib.utils as oneutils

# TODO Find better way to suppress trackback on error
sys.tracebacklimit = 0

subtool_list = {
    'compile': {
        'import': 'Convert given model to circle',
        'optimize': 'Optimize circle model',
        'quantize': 'Quantize circle model',
    },
    'package': {
        'pack': 'Package circle and metadata into nnpackage',
    },
    'backend': {
        'codegen': 'Code generation tool',
        'profile': 'Profile backend model file',
        'infer': 'Infer backend model file'
    },
}


def _call_driver(driver_name, options):
    dir_path = os.path.dirname(os.path.realpath(__file__))
    driver_path = os.path.join(dir_path, driver_name)
    cmd = [driver_path] + options
    oneutils.run(cmd)


def _check_subtool_exists():
    """verify given arguments"""
    subtool_keys = [n for k, v in subtool_list.items() for n in v.keys()]
    if len(sys.argv) > 1 and sys.argv[1] in subtool_keys:
        driver_name = 'one-' + sys.argv[1]
        options = sys.argv[2:]
        _call_driver(driver_name, options)
        sys.exit(0)


def _get_parser():
    onecc_usage = 'onecc [-h] [-v] [-C CONFIG] [-W WORKFLOW] [-O OPTIMIZATION] [COMMAND <args>]'
    onecc_desc = 'Run ONE driver via several commands or configuration file'
    parser = argparse.ArgumentParser(description=onecc_desc, usage=onecc_usage)

    oneutils.add_default_arg(parser)

    opt_name_list = oneutils.get_optimization_list(get_name=True)
    opt_name_list = ['-' + s for s in opt_name_list]
    if not opt_name_list:
        opt_help_message = '(No available optimization options)'
    else:
        opt_help_message = '(Available optimization options: ' + ', '.join(
            opt_name_list) + ')'
    opt_help_message = 'optimization name to use ' + opt_help_message
    parser.add_argument('-O', type=str, metavar='OPTIMIZATION', help=opt_help_message)

    parser.add_argument(
        '-W', '--workflow', type=str, metavar='WORKFLOW', help='run with workflow file')

    # just for help message
    compile_group = parser.add_argument_group('compile to circle model')
    for tool, desc in subtool_list['compile'].items():
        compile_group.add_argument(tool, action='store_true', help=desc)

    package_group = parser.add_argument_group('package circle model')
    for tool, desc in subtool_list['package'].items():
        package_group.add_argument(tool, action='store_true', help=desc)

    backend_group = parser.add_argument_group('run backend tools')
    for tool, desc in subtool_list['backend'].items():
        backend_group.add_argument(tool, action='store_true', help=desc)

    return parser


def _parse_arg(parser):
    args = parser.parse_args()
    # print version
    if args.version:
        oneutils.print_version_and_exit(__file__)

    return args


def _verify_arg(parser, args):
    """verify given arguments"""
    # check if required arguments is given
    if not oneutils.is_valid_attr(args, 'config') and not oneutils.is_valid_attr(
            args, 'workflow'):
        parser.error('-C/--config or -W/--workflow argument is required')
    # check if given optimization option exists
    opt_name_list = oneutils.get_optimization_list(get_name=True)
    opt_name_list = [oneutils.remove_prefix(s, 'O') for s in opt_name_list]
    if oneutils.is_valid_attr(args, 'O'):
        if ' ' in getattr(args, 'O'):
            parser.error('Not allowed to have space in the optimization name')
        if not getattr(args, 'O') in opt_name_list:
            parser.error('Invalid optimization option')


def main():
    # check if there is subtool argument
    # if true, it executes subtool with argv
    # NOTE:
    # Why call subtool directly without using Argparse?
    # Because if Argparse is used, options equivalent to onecc including
    # '--help', '-C' are processed directly onecc itself.
    # So options cannot be delivered to subtool.
    _check_subtool_exists()

    # parse arguments
    # since the configuration file path is required first,
    # parsing of the configuration file proceeds after this.
    parser = _get_parser()
    args = _parse_arg(parser)

    # verify arguments
    _verify_arg(parser, args)

    bin_dir = os.path.dirname(os.path.realpath(__file__))
    if oneutils.is_valid_attr(args, 'config'):
        runner = CfgRunner(args.config)
        runner.detect_import_drivers(bin_dir)
        if oneutils.is_valid_attr(args, 'O'):
            runner.add_opt(getattr(args, 'O'))
        runner.run(bin_dir)
    elif oneutils.is_valid_attr(args, 'workflow'):
        runner = WorkflowRunner(args.workflow)
        runner.run(bin_dir)


if __name__ == '__main__':
    oneutils.safemain(main, __file__)