#!/usr/bin/env python
#
# Generate FSL conda environment files from fsl-release.yml


import itertools  as it
import os.path    as op
import               os
import               re
import               sys

from fsl_ci.conda         import  get_channel_packages
from manifest_rules.utils import (load_release_info,
                                  generate_environment_file_name,
                                  generate_development_version_identifier)


def generate_package_list(release_info, platform, cudaver):
    """Return a list of packages that should be included in the environment
    file for the platform and CUDA version.
    """
    pkgs      = release_info['packages']
    platpkgs  = release_info.get(f'{platform}-packages', [])
    pkgs      = pkgs + platpkgs

    cudapat   = r'cuda-[\d]+\.[\d]+'
    cudapkgs  = [p for p in pkgs if re.search(cudapat, p) is not None]
    otherpkgs = [p for p in pkgs if p not in cudapkgs]

    if cudaver is None:
        return otherpkgs

    else:
        include  = f'cuda-{cudaver}'
        cudapkgs = [p for p in cudapkgs if include in p]

        return otherpkgs + cudapkgs


def need_internal_channel(release_info, packages):
    """Returns True if any of the given packages are sourced from the
    internal FSL conda channel.
    """

    username          = os.environ['FSLCONDA_USERNAME']
    password          = os.environ['FSLCONDA_PASSWORD']
    channel_url       = release_info['internal-channel']
    channel_url       = channel_url.replace('${FSLCONDA_USERNAME}:${FSLCONDA_PASSWORD}@', '')
    internal_packages = get_channel_packages(channel_url,
                                             username=username,
                                             password=password)
    internal_packages = list(internal_packages.keys())
    packages          = [p.split()[0] for p in packages]

    return any([p in internal_packages for p in packages])


def generate_environment(release_info, version, platform, cudaver, outfile):
    """Genereate an environment file for the platform and CUDA version.
    """
    channels = list(release_info['channels'])
    packages = generate_package_list(release_info, platform, cudaver)

    # Internal/dev release - add internal
    # channel to environment spec if necessary
    if need_internal_channel(release_info, packages):
        channels = [release_info['internal-channel']] + channels

    with open(outfile, 'wt') as f:
        f.write('name: FSL\n')
        f.write('channels:\n')
        for channel in channels:
            f.write(f' - {channel}\n')
        f.write('dependencies:\n')
        for package in packages:
            f.write(f' - {package}\n')


def generate_variants(release_info):
    """Generate a list of (platform, cudaver) pairs for which an environment
    file should be generated. For non-CUDA environments, cudaver will be None.
    """
    cudas     = [None] + list(release_info['cuda'])
    platforms = list(release_info['miniconda'].keys())
    variants  = list(it.product(platforms, cudas))

    # CUDA builds are only supported
    # on linux-64 at this time
    variants = [(plat, cuda) for plat, cuda in variants
                if not (('linux-64' not in plat) and cuda is not None)]

    return variants


def main():
    # Save generated environment files here
    release_file = op.abspath(sys.argv[1])
    outdir       = op.abspath(sys.argv[2])

    # Tags on fsl/conda/manifest denote a
    # public FSL release.  All other
    # commits denote an internal release.
    version      = os.environ.get('CI_COMMIT_TAG', None)
    release_info = load_release_info(release_file)

    # Generate environment files
    # for all platforms+CUDA versions
    variants = generate_variants(release_info)

    if not op.exists(outdir):
        os.mkdir(outdir)

    generated = []

    for platform, cuda in variants:
        filename = generate_environment_file_name(version, platform, cuda)
        filename = op.join(outdir, filename)
        generate_environment(release_info, version, platform, cuda, filename)
        generated.append((platform, cuda, op.basename(filename)))

    for platform, cuda, filename in generated:
        print(f'Platform {platform} [CUDA {cuda}]: {filename}')


if __name__ == '__main__':
    sys.exit(main())