#!/usr/bin/env python
#
# Generate FSL release manifest file from fsl-release.yml


import os.path      as     op
import                     os
import                     sys
import                     glob
import                     json
import urllib.parse as     urlparse
from   collections  import defaultdict

from manifest_rules.utils import (load_release_info,
                                  tempdir,
                                  download_file,
                                  sha256,
                                  parse_environment_file_name,
                                  generate_development_version_identifier)


def load_test_install_info(install_dir):
    """Returns information (number of output lines produced) about
    the test installations. Returns them in a dict.
    """
    inst_info  = defaultdict(dict)
    inst_files = glob.glob(op.join(install_dir, '*.install.txt'))

    for fname in inst_files:
        with open(fname, 'rt') as f:
            lines = f.readlines()
            info = lines[-1].strip()

        envfile, inst_lines, env_lines = info.split()
        platform        = parse_environment_file_name(envfile)[1]
        inst_lines      = int(inst_lines)
        env_lines       = int(env_lines)
        prev_inst_lines = inst_info['miniconda'].get('platform', 0)
        inst_lines      = max((inst_lines, prev_inst_lines))

        inst_info['miniconda'][platform] = inst_lines
        inst_info[envfile]               = env_lines

    return inst_info


def load_previous_manifest(release_info):
    """Downloads the previous official FSL manifest, returning it
    as a dict. Returns None if the manifest cannot be downloaded.
    """
    url = urlparse.urljoin(release_info['release-url'], 'manifest.json')
    with tempdir():

        try:
            download_file(url, 'manifest.json')
        except Exception:
            return None

        with open('manifest.json', 'rt') as f:
            lines = f.readlines()

    # Drop comments
    lines = [l for l in lines if not l.lstrip().startswith('//')]

    manifest = json.loads('\n'.join(lines))

    # Filter out any releases/builds for which
    # the environment file is not available
    versions = dict(manifest.get('versions', {}))
    for version, builds in versions.items():
        if version == 'latest':
            continue
        for i, build in enumerate(list(builds)):
            try:
                with tempdir():
                    download_file(build['environment'], 'env.yml')
            except Exception:
                print(f'WARNING: Build environment file '
                      f'is not available: {build}')
                builds.pop(i)

        if len(builds) == 0:
            versions.pop(version)
        else:
            versions[version] = builds

    latest = versions.get('latest', None)
    if (latest is not None) and (latest not in versions):
        print(f'WARNING: Build for latest FSL '
              f'version {latest} is not available')
        versions.pop('latest')

    manifest['versions'] = versions

    return manifest


def generate_installer_section(release_info):
    """Generates the "installer" manifest section, returning a dict. """
    version  = release_info['installer']
    url      = urlparse.urljoin(release_info['release-url'], 'fslinstaller.py')

    try:
        with tempdir():
            download_file(url, 'fslinstaller.py')
            checksum = sha256('fslinstaller.py')
    except Exception:
        print(f'WARNING: fslinstaller.py script not available at {url}!')
        checksum = ''

    return {'version' : version,
            'url'     : url,
            'sha256'  : checksum}


def generate_miniconda_section(release_info, install_info):
    """Generates the "miniconda" manifest section, returning a dict. """
    section = {}

    outputs = install_info.get('miniconda', {})

    for platform, url in release_info['miniconda'].items():
        output = outputs.get(platform, None)
        with tempdir():
            download_file(url, 'miniconda.sh')
            checksum = sha256('miniconda.sh')

            section[platform] = {
                'url'    : url,
                'sha256' : checksum,
            }
            if output is not None:
                section[platform]['output'] = str(output)

    return section


def generate_version_section(version,
                             include_current_release,
                             include_past_releases,
                             envdir,
                             release_info,
                             install_info):
    """Generates the "versions" manifest section, returning a dict.
    Does not add the "latest" entry.
    """

    section = {}

    if include_past_releases:
        prev_manifest = load_previous_manifest(release_info)
        if prev_manifest is not None:
            section = prev_manifest['versions']

    if not include_current_release:
        return section

    versions = []
    section  = {version : versions}

    for envfile in glob.glob(op.join(envdir, '*.yml')):

        checksum = sha256(envfile)
        envfile  = op.basename(envfile)
        output   = install_info.get(envfile, None)
        url      = urlparse.urljoin(release_info['release-url'], envfile)
        version, platform, cuda = parse_environment_file_name(envfile)

        build = {
            'platform'      : platform,
            'environment'   : url,
            'sha256'        : checksum,
            'base_packages' : release_info['base-packages']
        }
        if cuda is not None:
            build['cuda'] = cuda

        versions.append(build)
        if output is not None:
            versions[-1]['output'] = {'install' : str(output)}

    return section


def main():
    # This script is run in the following scenarios:
    #   - New public FSL release - generate a new official
    #     manifest, describing the new release, and all past
    #     public releases
    #   - New internal/development release - generate a
    #     development manifest describing just the development
    #     release
    #   - Update to official manifest - re-generate the
    #     official manifest, describing all past public
    #     releases
    # We use the "official" and "add_current_release"
    # variables to figure out the scenario we are in

    release_file = op.abspath(sys.argv[1])

    # See .gitlab-ci.yml - we are passed "true" or "false"
    # to determine whether we are to generate an official
    # or development manifest
    official = sys.argv[2] == 'true'

    # Directory to save the generated manifest file
    outdir = op.abspath(sys.argv[3])

    # Directory containing the environment files
    # to be added to/described by the manifest
    envdir = op.abspath(sys.argv[4])

    # Directory containing outputs of the test
    # installation
    installdir = op.abspath(sys.argv[5])

    # FSL version identifier
    version = os.environ.get('CI_COMMIT_TAG', None)

    # Add information about the current release
    # (tag or commit) to the generated manifest?
    # We do this for new official releases, or
    # for development releases
    add_current_release = (version is not None) or (not official)

    # The "official" variable controls whether we
    # are generating the official "manifest.json",
    # which typically includes information about
    # past official FSL releases (which is done by
    # downloading the currently available official
    # release manifest).  This can be overridden by
    # setting the NO_PAST_RELEASES variable.
    add_past_releases = official and ('NO_PAST_RELEASES' not in os.environ)

    # Generate dev version identnfier
    if version is None:
        version = generate_development_version_identifier()

    # Load fsl-release.yml, and load
    # outputs of test install stage
    release_info = load_release_info(release_file)
    install_info = load_test_install_info(installdir)

    # Generate the manifest
    manifest = {
        'installer' : generate_installer_section(release_info),
        'miniconda' : generate_miniconda_section(release_info, install_info),
        'versions'  : generate_version_section(version,
                                               add_current_release,
                                               add_past_releases,
                                               envdir,
                                               release_info,
                                               install_info),
    }

    # Set the "latest" field for for
    # new public or developmenmt releases
    if add_current_release:
        manifest['versions']['latest'] = version

    if official:  fname = 'manifest.json'
    else:         fname = f'manifest-{version}.json'

    manifest = json.dumps(manifest, indent=4, sort_keys=True)

    with open(op.join(outdir, fname), 'wt') as f:
        f.write(manifest)

    print(f'Generated manifest file {fname}:')
    print(manifest)


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