#!python
#
# 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.0OA
#
# Authors:
# - Wen Guan, <wen.guan@cern.ch>, 2024 - 2025


"""
Run workflow.
"""

from __future__ import print_function

import argparse

# import argcomplete
import logging
import os
import sys
import time
import traceback


from idds.common.utils import setup_logging
from idds.common.version import release_version
from idds.prompt.transformer import Transformer

setup_logging(__name__, stream=sys.stdout)


def run_prompt(args):
    logging.info(f"run_prompt called with args: {str(args)}")

    transformer = Transformer(run_id=args.run_id, workdir=args.workdir, namespace=args.namespace, idle_timeout=args.idle_timeout)
    logging.info("Initializing transformer brokers")
    ret = transformer.run()
    logging.info("Transformer run returned: %s" % str(ret))
    return ret


def get_parser():
    """
    Return the argparse parser.
    """
    oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)

    # common items
    oparser.add_argument(
        "--version", action="version", version="%(prog)s " + release_version
    )
    oparser.add_argument(
        "--verbose",
        "-v",
        default=False,
        action="store_true",
        help="Print more verbose output.",
    )
    oparser.add_argument(
        "--run_id", type=int, required=True, help="Run ID for the prompt run."
    )
    oparser.add_argument(
        "--workdir", type=str, required=False, default=None, help="Working directory for the prompt run."
    )
    oparser.add_argument(
        "--namespace", type=str, required=False, default=None, help="Namespace for the prompt run."
    )

    oparser.add_argument(
        "--idle_timeout",
        type=int,
        default=600,
        help="Idle timeout in seconds to stop the transformer if no messages are received.",
    )

    return oparser


if __name__ == "__main__":
    arguments = sys.argv[1:]

    oparser = get_parser()
    # argcomplete.autocomplete(oparser)

    args, unknown = oparser.parse_known_args(arguments)

    try:
        if args.verbose:
            logging.getLogger().setLevel(logging.DEBUG)
        start_time = time.time()
        logging.info("Starting prompt with arguments: {0}".format(arguments))
        logging.info("Starting at: %-0.4f" % start_time)

        exit_code = run_prompt(args)

        end_time = time.time()
        logging.info("Finished at: %-0.4f" % end_time)
        logging.info("Completed in %-0.4f sec." % (end_time - start_time))
        sys.exit(exit_code)
    except Exception as error:
        logging.error("Strange error: {0}".format(error))
        logging.error(traceback.format_exc())
        end_time = time.time()
        logging.info("Failed at: %-0.4f" % end_time)
        logging.info("Failed in %-0.4f sec." % (end_time - start_time))
        sys.exit(-1)
