#!python
import argparse
import logging
import multiprocessing as mp
import os
import sys
import threading
import time
import traceback

import coloredlogs
from alibuild_helpers.utilities import detectArch

import flp2rpm.config as config
from flp2rpm.Fpm import Fpm
from flp2rpm.Module import Module
from flp2rpm.Package import Package
from flp2rpm.Repo import Repo
from flp2rpm.S3 import S3


def configure_logging():
    coloredlogs.install(level=config.log_level)


def set_config_args(args):
    """Sets the configuration variables as parsed from argparse Namespace or dict-like object."""
    # Helper to get attribute or default if not present
    def get_arg(attr, default=None):
        return getattr(args, attr, default) if hasattr(args, attr) else default

    config.ali_prefix = get_arg('ali_prefix', config.ali_prefix)
    logging.info(f"Using aliBuild work directory: {config.ali_prefix}")

    config.release_tag = get_arg('release_tag', config.release_tag)

    config.dry_run = get_arg('dry_run', config.dry_run)
    logging.info(f"Starting dry run: {config.dry_run}")

    config.fail_on_error = get_arg('fail_on_error', config.fail_on_error)
    logging.info(f"Interrupting generation on first RPM generation failure")

    config.skip_deps = get_arg('skip_deps', config.skip_deps)
    logging.info(f"Skipping dependencies: {config.skip_deps}")

    config.devel = get_arg('devel', config.devel)
    logging.info(f"Generating devel dependencies: {config.devel}")

    config.architecture = get_arg('architecture', detectArch())
    logging.info(f"Detected aliBuild architecture: {config.architecture}")

    config.target_rpm_dir = get_arg('target_rpm_dir', config.target_rpm_dir)
    config.tag_version = get_arg('tag_version', getattr(config, 'tag_version', False))
    config.workers = get_arg('workers', getattr(config, 'workers', 15))

    config.target_rpm_dir = os.path.join(config.target_rpm_dir, config.release_tag, 'o2')
    logging.info(f"Using output directory: {config.target_rpm_dir}")


def process_deco(func):
    """ Decorator for generate_rpm: recursively calls process_wrap on the dependencies of the package
        and call the generate_rpm on all of them.
        process_deco.processed keeps a list of everything that has already been built.
        System dependencies are skipped.
        It creates a thread for each rpm generation.
        DEPRECATED
    """
    process_deco.processed = []

    def process_wrap(name, version):
        logging.debug(f"Processing {name} {version}")
        if version == 'from_system':
            logging.debug("  from_system")
            return

        package = Package(name, version)
        deps = package.deps_with_versions()
        devel_deps = package.get_devel_deps_with_versions()
        extra_deps = package.get_extra_deps()

        if not config.skip_deps:
            for dep_name, dep_version in deps.items():
                process_wrap(dep_name, dep_version)
            for dev_name, dev_version in devel_deps.items():
                process_wrap(dev_name, dev_version)

        if name not in process_deco.processed:
            process_deco.processed.append(name)
            process_wrap.threads.append(func(name, version, deps, devel_deps, extra_deps))
            return process_wrap.threads

    process_wrap.threads = []
    return process_wrap


@process_deco
def generate_rpm(name, version, deps=None, devel_deps=None, extra_deps=None) -> threading.Thread:
    """DEPRECATED"""
    fpm = Fpm()
    t = threading.Thread(target=fpm.run, args=(name, version, deps, devel_deps, extra_deps))
    t.start()
    return t

def collect_packages_list(name, version) -> set[tuple[bool, str]]:
    """ Recursively build the set of all the packages to build for this package and its dependencies. """
    result = set()

    if version == 'from_system':
        logging.debug("  from_system")
        return result

    package = Package(name, version)
    deps = package.deps_with_versions()
    devel_deps = package.get_devel_deps_with_versions()

    if not config.skip_deps:
        for dep_name, dep_version in deps.items(): # runtime deps
            result |= collect_packages_list(dep_name, dep_version)
        for dev_name, dev_version in devel_deps.items(): # devel deps
            result |= collect_packages_list(dev_name, dev_version)
    result.add((name, version))
    return result

def generate_package(name, version, failures_list):
    """ Prepare the RPM for the package. In case of failure, it is added to the failures_list """
    package = Package(name, version)
    deps = package.deps_with_versions()
    devel_deps = package.get_devel_deps_with_versions()
    extra_deps = package.get_extra_deps()
    try:
        fpm = Fpm()
        status = fpm.run(name, version, deps, devel_deps, extra_deps)
    except:
        logging.error(f"Exception in generate_package: {traceback.format_exc()}")
        status = False
    if not status:
        failures_list.append((name, version))

def generate_rpm_with_pool():
    '''
    Build a list of all the packages to build and then use a Pool of processes to do the work.
    As we cannot be sure we won't run out of memory with large packges we retry for the ones that failed.
    '''
    packages_set = collect_packages_list(args.package, args.version)
    logging.debug(f"List of packages to build : {packages_set}")

    def error_callback(e):
        print(f"Error: {e}", file=sys.stderr)
        # This will terminate the pool and the main process
        pool.terminate()
        pool.join()
        sys.exit(1)

    # first attempt
    pool = mp.Pool(processes=int(args.workers))
    print(f"Using {args.workers} workers to generate RPMs")
    manager = mp.Manager()
    failures_list = manager.list() # list to collect the packages that failed to build
    [pool.apply_async(generate_package, (name, version, failures_list), error_callback=error_callback) for name, version in packages_set]
    pool.close()
    pool.join()

    # retries
    if len(failures_list) != 0:
        pool = mp.Pool(processes=int(args.workers))
        logging.info(f"List of failed packages {failures_list}. \nRetry now")
        [pool.apply_async(generate_package, (name, version, failures_list)) for name, version in failures_list]
        pool.close()
        pool.join()


def main(args):
    start_main = time.time()
    if 'log_level' in args:
        config.log_level = args.log_level
    configure_logging()
    set_config_args(args)
    if args.command == 'generate':
        if args.use_pool:
            generate_rpm_with_pool()
        else: # this is the old way of doing it. We keep it for the time being as a fallback.
            thread_list = generate_rpm(args.package, args.version)
            logging.info('Waiting for RPM generation to finish...')
            for t in thread_list: # BvH: this makes no sense, we have a single thread
                t.join()
        logging.info('RPM generation completed in %.1f seconds' % (time.time() - start_main))
    if args.command == 'validate':
        Repo.validate_rpms()
        logging.info('RPM validation completed in %.1f seconds' % (time.time() - start_main))
    if args.command == 'repo':
        Repo.create()
        logging.info('Repo create completed in %.1f seconds' % (time.time() - start_main))
    if args.command == 'sync' and not args.pull:
        S3().push_rpms(args.delete_removed)
        logging.info('Sync push completed in %.1f seconds' % (time.time() - start_main))
    if args.command == 'sync' and args.pull:
        S3().pull_rpms()
        logging.info('Sync pull completed in %.1f seconds' % (time.time() - start_main))
    if args.command == 'module':
        print('\n'.join(Module(args.avail).versions(False)))
    if args.command == 's3-copy':
        S3().copy(args.from_tag, args.to, args.delete_removed)
        logging.info('Copy completed in %.1f seconds' % (time.time() - start_main))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--dry-run', help='do a dry run, skipping fpm execution', action="store_true", required=False, default=argparse.SUPPRESS)
    parser.add_argument('--target-rpm-dir', help='path to store RPMs in (=' + config.target_rpm_dir + ' by default)', required=False, default=argparse.SUPPRESS)
    parser.add_argument('--release-tag', help='Release tag, this is mostly to provide correct dir structure', required=False, default=argparse.SUPPRESS)
    parser.add_argument('--architecture', help='OS architecture', required=False, default=argparse.SUPPRESS)
    parser.add_argument('--log-level', help='Set log level (DEBUG, INFO, WARN, ERROR)', required=False, default=argparse.SUPPRESS)
    subparser = parser.add_subparsers(help='Available commands:', dest='command')
    # generate command
    parser_generate = subparser.add_parser('generate', help='Generate RPMs')
    parser_generate.add_argument('--package', help='package name (as recipe name in alidist)', required=True)
    parser_generate.add_argument('--version', help='package version (as in modulefile: X.Y.Z-A)', required=True)
    parser_generate.add_argument('--ali-prefix', help='path to alibuild dir', required=False, default=argparse.SUPPRESS)
    parser_generate.add_argument('--skip-deps', help='Generate single RPM without dependencies', action="store_true", required=False, default=argparse.SUPPRESS)
    parser_generate.add_argument('--devel', help='Generate also devel RPM', action="store_true", required=False, default=argparse.SUPPRESS)
    parser_generate.add_argument('--tag-version', help='Add release tag to RPM version', action="store_true", required=False, default=argparse.SUPPRESS)
    parser_generate.add_argument('--workers', help='Number of parallel worker to generate RPMs (only work with --use-pool)', required=False, default=15)
    parser_generate.add_argument('--use-pool', help='Use a pool of worker and not an infinite number of threads', required=False, action="store_true")
    # sync
    parser_sync = subparser.add_parser('sync', help='Sync RPMs to S3')
    parser_sync.add_argument('--delete-removed', help='Remove locally deleted files from S3', action="store_true", required=False, default=False)
    parser_sync.add_argument('--pull', help='Pulls instead of pushing', action="store_true", required=False, default=False)
    # s3
    parser_s3 = subparser.add_parser('s3-copy', help='Copy files between <folders> on S3')
    parser_s3.add_argument('--from', help='Origin as release tag (eg. dev)', required=True, dest='from_tag')
    parser_s3.add_argument('--to', help='Destination as release tag (eg. flp-suite-v1.0)', required=True)
    parser_s3.add_argument('--delete-removed', help='Remove destination files no longer available in origin', action="store_true", required=False, default=False)
  
    # validate
    parser_validate = subparser.add_parser('validate', help='Validate RPMs')
    # repo
    parser_repo = subparser.add_parser('repo', help='Create YUM repo')
    # module
    parser_module = subparser.add_parser('module', help='Operate on modulefiles')
    parser_module.add_argument('--avail', help='Display available modules for given prefix', required=True)
    parser_module.add_argument('--ali-prefix', help='path to alibuild dir', required=False, default=argparse.SUPPRESS)
    args = parser.parse_args()
    main(args)
    sys.exit(0)

