#!/usr/bin/env python3
import argparse
import sys
import os
import logging
from modules import opt_module
from utils import log
from utils import util


def arg_parse():
    desc = 'Environment variables:\n    - COMPILER_BIN_DIR: specify compiler bin dir\n\n'
    desc += 'Sub-modules:\n    - image_cut_search: 720|530|730|630|540\n'
    parser = argparse.ArgumentParser(description=desc, formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("platform", help="HW platform <520|720|530|730|630|540>")
    parser.add_argument("model", help="onnx or bie model path")
    parser.add_argument("config", help="compiler base config. config.json")
    parser.add_argument('output_dir', help='log output directory')
    parser.add_argument("-d", "--debug", action="store_true", default=False, help="debug mode. Show more debug message")

    return parser.parse_args()


if __name__ == "__main__":

    # get opt wrapper directory
    opt_dir = os.path.dirname(os.path.abspath(__file__))

    # get arguments
    args = arg_parse()
    platform = args.platform
    model = args.model
    base_config = args.config
    output_dir = args.output_dir
    log_level = logging.DEBUG if args.debug else logging.INFO

    # create logger and output directory
    log_path = output_dir if output_dir else '__magic_tmp__'
    remove_log = not output_dir
    log = log.Log(log_path, remove_log)
    log.copy_files(model, base_config)
    log.logger.setLevel(log_level)

    if output_dir and not util.is_file_name_valid(output_dir):
        log.logger.error('Please give a valid -o/OUTPUT_DIR. "{}" is invalid'.format(output_dir))
        sys.exit(-1)

    # set compiler binary path to environment. the env is used in sub-module
    compiler_bin = os.environ['COMPILER_BIN'] = util.get_compiler_bin(opt_dir)
    if not util.check_file_exist(compiler_bin):
        log.logger.error('compiler binary [path: {}] not exist!'.format(compiler_bin))
        sys.exit(-1)


    # run optimize modules
    opt_modules = opt_module.OptModules(platform, model, base_config, opt_dir, log)
    log.logger.info('==== Information ====')
    log.logger.info('opt_compile     : {}'.format(' '.join(sys.argv)))
    log.logger.info('platform        : {}'.format(platform))
    log.logger.info('model           : {}'.format(model))
    log.logger.info('base config     : {}'.format(base_config))
    log.logger.info('sub-modules     : {}'.format(opt_modules.modules))
    log.logger.info('log directory   : {}'.format(log_path))
    log.logger.info('compiler bin    : {}'.format(compiler_bin))
    log.logger.info('=====================')
    opt_modules.run()

    # TODO: find best order for opt_modules