#!/usr/bin/python3
import argparse
import yaml
import os
import os.path as path
import shutil
import platform
import urllib.request
import subprocess
import collections
import sys
from packaging.version import Version
from typing import Optional, List, Dict

def get_mamba_bin(base_path):
    """
    Return the path to the mamba binary
    :param base_path: base path where mamba binary resides
    :return: path to mamba binary
    """
    return path.join(base_path,"bin","mamba")

def get_conda_bin(base_path: str) -> str:
    """
    Return path to conda binary
    """
    return path.join(base_path,"bin","conda")

def download_installer(installer_path):
    """
    Download the mamba installer

    :param installer_path:
    :return: path_to_installer
    """
    uname = platform.uname()
    name = f"Miniforge3-{uname.system}-{uname.machine}.sh"
    url = f"https://github.com/conda-forge/miniforge/releases/latest/download/{name}"

    local_installer = path.join(installer_path, name)

    print(f"Downloading {url} to {local_installer}")

    with urllib.request.urlopen(url) as u:
        with open(local_installer, "wb") as f:
            f.write(u.read())

    os.chmod(local_installer, 0o766)

    return local_installer


def install_conda(installer_path: str, base_path: str):
    """
    Install conda (or mamba more likely) if it's not already installed.
    :param installer_path: Where to download the installer.  Will be created if missing.
    :param base_path: where to install conda.  Will be created if missing.
    """
    install_flag = path.join(base_path, 'installed.txt')
    if path.exists(install_flag):
        print("Mamba already installed. Skipping installation.")
    else:
        print("Mamba not installed.  Installing.")

        # delete the base path if it already exists.
        # this probably means we had a bad installation
        if path.isdir(base_path):
            print("Found a partial installation.  Deleting.")
            shutil.rmtree(base_path)

        os.makedirs(installer_path, exist_ok=True)

        installer = download_installer(installer_path)
        print(f"running Miniforge installer, installing to {base_path}")
        result = subprocess.run([installer, "-b", "-p", base_path],
                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8')
        if result.returncode == 0:
            print("Miniforge installation successful")
        else:
            print("Miniforge installation failed")
            print("output was: ")
            print(result.stdout)
            raise Exception("Miniforge installation failed")

        mamba_exec = get_mamba_bin(base_path)

        print("installing conda-build")
        result = subprocess.run([mamba_exec, "install", "-q", "-y", "conda-build"],
                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8')

        if result.returncode == 0:
            print("conda-build installation successful")
        else:
            print("conda-build installation failed")
            print("output was:")
            print(result.stdout)
            raise Exception("conda-build installation failed")

        # create installation flag
        with open(install_flag, "wt"):
            pass


def read_config(config_path):
    """
    Read in the desired conda configuration from a yaml file.
    :param config_path:
    :return: The configuration as a regular python data structure
    """
    with open(config_path, "rt") as f:
        config = yaml.safe_load(f)
    return config


def install_environments(base_path, forge_url, log_file, environments, noop):
    """
    Install conda environments.  Will skip any that are already installed.
    :param environments: A list of environments
    :param forge_url: A url prefix used to find environment definitions
    :param log_file: an open file handle to write the output of the conda command
    :return: set of newly installed environments
    """
    error_code = 0
    noop_text = noop and " [NOOP]" or ""
    installed = []
    mamba_bin = get_mamba_bin(base_path)
    print("checking for new environments")
    for env in environments:
        env_path = path.join(base_path,"envs",env)
        if path.isdir(env_path):
            print(f"{env} already installed")
        else:
            print(f"installing {env} environment{noop_text}")
            url = f"{forge_url}/{env}.yaml"
            if not noop:
                temp_file = path.join(f"/tmp/{env}.yaml")
                result = subprocess.run(["/usr/bin/wget", url, '-O', temp_file])
                if result.returncode != 0:
                    print(f"Failed to fetch environment {env}")
                else:
                    # environ = os.environ.copy()
                    # environ["CONDA_OVERRIDE_ARCHSPEC"] = "x86_64"
                    result = subprocess.run([mamba_bin, 'env', "create", '--quiet', '-y', '--file', temp_file],
                                   stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8',
                                            )
                                   # env=environ)
                    if result.returncode == 0:
                        print(f"{env} installed")
                        installed.append(env)
                    else:
                        error_code = 1
                        print(f"{env} install failed")
                    if log_file is not None:
                        log_file.write(result.stdout)
    return error_code, set(installed)


def update_environments(base_path, forge_url, log_file, environments, noop):
    """
    Update selected conda environments.
    :param environments: A list of environments
    :param forge_url: A url prefix used to find environment definitions
    :param log_file: an open file handle to write the output of the conda command
    :return:
    """
    noop_text = noop and " [NOOP]" or ""
    mamba_bin = get_mamba_bin(base_path)
    print("Updating existing environments")
    error_code = 0
    for env in environments:
        print(f"updating {env} environment{noop_text}")
        url = f"{forge_url}/{env}.yaml"
        if not noop:
            temp_file = path.join(f"/tmp/{env}.yaml")
            result = subprocess.run(["/usr/bin/wget", url, '-O', temp_file])
            if result.returncode != 0:
                print(f"Failed to fetch environment {env}")
            else:
                result = subprocess.run([mamba_bin, "env", "update", "--yes", "--file", temp_file],
                                        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8')
                if result.returncode == 0:
                    print(f"{env} updated")
                else:
                    error_code = 1
                    print(f"{env} update failed")
                if log_file is not None:
                    log_file.write(result.stdout)
    return error_code


def remove_environments(base_path, envs, noop):
    error_code = 0
    noop_text = noop and " [NOOP]" or ""
    for env in envs:
        env_path = path.join(base_path, "envs", env)
        try:
            if path.islink(env_path):
                print(f"Deleting {env}{noop_text}")
                if not noop:
                    os.remove(env_path)
            elif path.isdir(env_path):
                print(f"Deleting {env}{noop_text}")
                if not noop:
                    shutil.rmtree(env_path)
            else:
                pass # not an environment
        except Exception as e:
            error_code = 1
            print(f"could not remove environment {env}: " + str(e))
    return error_code

def find_unreferenced_environments(base_path, wanted_envs):
    """
    Remove any file (directory, or symlink) not referenced in installed or linked environments
    :param base_path:
    :param installed_envs:
    :return: A list of deleted environments
    """

    print("Preparing to remove unneeded environments")
    installed_envs = set(os.listdir(path.join(base_path,"envs")))
    delete_envs = installed_envs - set(wanted_envs)
    return delete_envs

def clean_cache(base_path, noop):
    mamba_bin = get_mamba_bin(base_path)
    print("Removing unused conda packages, tarfiles from {base_path}")
    if not noop:
        result = subprocess.run([mamba_bin, 'clean', '-y', '-p', '-t'],
                       stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8')
        if result.returncode == 0:
            print(f"{base_path} caches cleaned")
        else:
            print(f"{base_path} cache clean failed:")
            print(result.stdout)
            print(result.stderr)

def create_links(base_path, links, noop):
    """
    Create symlink environments
    :param base_path: base path to conda installation
    :param links: tuple of (env_name, env_path)
    :return:
    """
    error_code = 0
    print("Creating symlinks to external environments")
    noop_text = noop and " [NOOP]" or ""
    for link in links:
        env_path = link[1]
        env_name = link[0]
        link_path = path.join(base_path, 'envs', env_name)
        if not (path.exists(link_path) or path.islink(link_path)):
            try:
                if not noop:
                    os.symlink(env_path, link_path)
            except Exception as e:
                error_code = 1
                print(f"could not create environment {env_name} as symlink to {env_path}: " + str(e) + noop_text)
            else:
                print(f"Created {env_name} as link to {env_path} {noop_text}")
        elif path.islink(link_path):
            link_targ_path = path.realpath(link_path)
            if not (path.exists(link_targ_path) and path.exists(env_path) and path.samefile(link_targ_path, env_path)):
                print(f"{env_name} already exists, but is pointing at a different path.  Correcting. {noop_text}")
                try:
                    if not noop:
                        os.remove(link_path)
                except Exception as e:
                    error_code = 1
                    print(f"could not remove existing link for {env_name}: " + str(e))
                else:
                    try:
                        if not noop:
                            os.symlink(env_path, link_path)
                    except Exception as e:
                        error_code = 1
                        print(f"could not create environment {env_name} as symlink to {env_path}: " + str(e))
                    else:
                        print(f"Created {env_name} as link to {env_path} {noop_text}")
            else:
                print(f"{env_name} exists and is pointing at the correct path")
        else:
            print(f"Couldn't create link at {link_path} because the file already exists")
    return error_code

def check_links(links):
    """
    Check desired symlinks for various problems

    Check for duplicates (Warning)
    Duplicates pointing at different paths (Error)
    Links to paths that aren't absolute (Error)
    Links to paths that don't exist (Error)

    :param links: A list of link pairs with (environment_name, link_path)
    :return: If no errors, return, otherwise throw exception
    """
    duplicates = set([])
    disagree_dups = {}
    relative = []   # simple links to a sister environment are ok
    dont_exist = []
    targets = {}
    same_target = set([])

    record = {}
    for link in links:
        env_name = link[0]
        env_path = link[1]

        # duplicates
        if env_name in record:
            prev_path = record[env_name]
            if path.exists(prev_path) and path.exists(env_path) and path.samefile(prev_path, env_path):
                duplicates.add(link)
            else:
                if env_name in disagree_dups:
                    disagree_dups[env_name].add(env_path)
                else:
                    disagree_dups[env_name] = set([prev_path, env_path])
        else:
            record[env_name] = env_path

        if not path.isabs(env_path):
            if not (path.basename(env_path) == env_path):
                # simple relative links to environments in the same directory are OK.
                relative.append(link)

        if not path.exists(env_path):
            # only error for absolute paths.  simple paths we don't care if it exists yet
            if (path.basename(env_path) == env_path):
                print(f"*** Warning: creating link to environment {env_path} but it does not exist.")
            else:
                dont_exist.append(link)

        real_path = path.realpath(env_path)

        if real_path in targets:
            same_target.add(real_path)
            targets[real_path].add(env_name)
        else:
            targets[real_path] = set([env_name])

    if len(duplicates) > 0:
        dups = [link[0] for link in duplicates]
        print("*** Warning: the following environments have more than one link requested: " + ", ".join(dups))

    for realpath in same_target:
        print(f"*** Warning: {realpath} has more than one environment pointing at it: " + ", ".join(targets[realpath]))

    error = False

    for env, paths in disagree_dups.items():
        error = True
        print(f"*** Error: requested multiple different target paths for {env}: " + ":".join(paths))

    for link in relative:
        error = True
        env = link[0]
        env_path = link[1]
        print(f"*** Error: target for {env} must be an absolute path, but relative path given: {env_path}")

    for link in dont_exist:
        error = True
        env = link[0]
        env_path = link[1]
        print(f"*** Error: requested target for {env} => {env_path} does not exist.")

    if error:
        raise Exception("One or more errors were found in requested links for environments")


def read_user_links(config_files):
    """
    Open up and read user config_files, return a list of symlinks found there.

    Yaml files are expected to be:

    envname: target_path

    :param config_files:  A list of config files to read in.
    :return:
    """
    links = []
    for conffile in config_files:
        print(f"reading user config file: {conffile}")
        with open(conffile, "rt") as f:
            conf = yaml.safe_load(f)
        for env_name, env_path in conf.items():
            if isinstance(env_name, str) and isinstance(env_path, str):
                links.append((env_name, env_path))
            else:
                print(f"*** Warning: unrecognized configuration {env_name} in {conffile}")
    return links

def get_version(command: str) -> Optional[Version]:
    """
    Get the version of either mamba or conda
    This is always the last word of the first line of the output.
    """
    result = subprocess.run([command, '--version'], stdout=subprocess.PIPE, encoding='utf-8')
    if result.returncode == 0:
        return Version(result.stdout.split("\n")[0].split()[-1])
    else:
        return None

class PackageSpec(object):
    """
    Identify a package and desired version range for it
    """
    def __init__(self, name: str, version_at_least: Optional[Version], version_less_than: Optional[Version]):
        self.name = name
        self.version_at_least = version_at_least
        self.version_less_than = version_less_than

    def __str__(self):
        """Produce a conda-style package specification with contraints.update_mamba_conda

        "None" version constraints are ignored.

        The output of this function is suitable as an argument to the `conda install` command.
        """
        version_strs = []
        if self.version_at_least is not None:
            version_strs.append(f">={self.version_at_least}")
        if self.version_less_than is not None:
            version_strs.append(f"<{self.version_less_than}")
        return self.name + ",".join(version_strs)

def update_base_packages(base_path: str, package_specs: List[PackageSpec]) -> int:
    """
    Update a list of packages complete with version range, typically conda or mamba, to match the version ranges
    given in the package specs.
    None values for versions ranges are ignored.
    """
    error_code = 0
    if len(package_specs) == 0:
        print("Not updating base packages")
        return 0

    print("updating these base packages: " + " ".join([n.name for n in package_specs]))
    command = [f"{base_path}/../update_package", base_path] + [str(p) for p in package_specs]
    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8')

    if result.returncode == 0:
        print("packages were updated")
    else:
        error_code = 1
        print("failed to update packages")
        print(f"stdout={result.stdout}")
    return error_code


def check_update_base_package(base_path: str, bin: str, package_spec: PackageSpec) -> Optional[bool]:
    """return True if the package needs to be updated. Return None if there's an error.
    """
    version = get_version(bin)
    if version is None:
        return None
    v_start = package_spec.version_at_least
    v_end = package_spec.version_less_than
    if (v_start is not None and version < v_start) or (v_end is not None and version >= v_end):
        print(f"{bin} version = {version}, which is not in target range [{v_start}, {v_end})")
        return True
    else:
        print(f"{bin} version = {version}, which is in target range [{v_start}, {v_end})")
        return False


def check_update_conda(base_path: str, package_spec: PackageSpec) -> Optional[bool]:
    """
    Update conda to the correct range, if needed
    """
    bin = get_conda_bin(base_path)
    return check_update_base_package(base_path, bin, package_spec)


def check_update_mamba(base_path: str, package_spec: PackageSpec) -> Optional[bool]:
    """
    Update mamba to the correct range, if needed
    """
    bin = get_mamba_bin(base_path)
    return check_update_base_package(base_path, bin, package_spec)

def convert_version_string(version_str: str) -> Optional[Version]:
    """Convert a version string into a Version object or None if the string is empty """
    if not version_str:
        return None
    return Version(version_str)

def update_mamba_conda(config: Dict, log_file, base_path: str, noop: bool) -> int:
    """update mamba and conda if needed"""
    mamba_spec = PackageSpec("mamba", convert_version_string(config['mamba_version_at_least']) or None, convert_version_string(config['mamba_version_less_than']) or None)
    conda_spec = PackageSpec("conda", convert_version_string(config['conda_version_at_least']) or None, convert_version_string(config['conda_version_less_than']) or None)

    to_upgrade: List[PackageSpec] = []

    error_code = 0

    r = check_update_conda(base_path, conda_spec)
    if r is None:
        error_code = 1
    elif r:
        to_upgrade.append(conda_spec)
    r = check_update_mamba(base_path, mamba_spec)
    if r is None:
        error_code = 1
    elif r:
        to_upgrade.append(mamba_spec)

    if noop:
        if len(to_upgrade) > 0:
            print("[NOOP] would have upgraded these base packages: " + ".".join([p.name for p in to_upgrade]))
        else:
            print("[NOOP] would not have upgraded any base packages")
    else:
        error_code = update_base_packages(base_path, to_upgrade) or error_code

    # make sure to update shell scripts
    command = f"{base_path}/../shell_init"
    result = subprocess.run([command, base_path],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding='utf-8')
    if result.returncode == 0:
        print("Updated mamba shell")
    else:
        error_code = result.returncode
        print("Update of mamba shell failed")
        if log_file is not None:
            log_file.write("# output for updating mamba shell\n")
            log_file.write(result.stdout)
    return error_code


def run(config_path, noop):

    error_code = 0
    config = read_config(config_path)
    print(config)

    conda_home = path.abspath(config['home'])

    conda_log = path.join(conda_home, "conda.log")

    environs = config['environs']

    remove_missing = config['remove_missing']

    forge_url = config['forge_url']

    print(f"conda_home={conda_home}")

    conda_base = path.join(conda_home, 'base')
    conda_install = path.join(conda_home, 'install')

    install_conda(conda_install, conda_base)

    for env in environs:
        if environs[env] is None:
            environs[env] = {}

    install_envs_list = [env for env, flags in environs.items() if ('ensure' not in flags or flags['ensure'] == "present") and 'target' not in flags]
    remove_envs_list = [env for env, flags in environs.items() if 'ensure' in flags and flags['ensure'] == "absent"]
    update_envs_list = [env for env, flags in environs.items() if 'update' in flags and flags['update'] and env in install_envs_list]

    # get links from yaml
    links_list = [(env, flags['target']) for env, flags in environs.items() if 'target' in flags]

    # add user links
    user_configs = ('user_configs' in config) and config['user_configs'] or []
    links_list += read_user_links(user_configs)

    link_envs_list = [link[0] for link in links_list]

    # look for duplicates
    install_dups = [item for item, count in collections.Counter(install_envs_list).items() if count > 1]
    remove_dups = [item for item, count in collections.Counter(remove_envs_list).items() if count > 1]

    if len(install_dups) > 1:
        print("*** WARNING, these environments are configured to be installed multiple times: " + ", ".join(install_dups))
    if len(remove_dups) > 1:
        print("*** WARNING, these environments are configured to be removed multiple times: " + ", ".join(remove_dups))

    check_links(links_list)

    # look for inconsistencies
    install_envs = set(install_envs_list)
    update_envs = set(update_envs_list)
    remove_envs = set(remove_envs_list)
    link_envs = set(link_envs_list)

    # process environments

    install_remove = install_envs.intersection(remove_envs)
    if len(install_remove) > 0:
        print("*** ERROR, these environments are configured to be both installed and removed: ", ", ".join(install_remove))
        raise Exception("Environments configured to be both installed and removed")

    install_link = install_envs.intersection(remove_envs)
    if len(install_link) > 0:
        print("*** ERROR, these environments are configured to be both installed and symmlinked: ", ", ".join(install_remove))
        raise Exception("Environments configured to be both installed and symmlinked")

    remove_link = link_envs.intersection(remove_envs)
    if len(remove_link) > 0:
        print("*** ERROR, these environments are configured to be both removed and symmlinked: ", ", ".join(install_remove))
        raise Exception("Environments configured to be both removed and symmlinked")

    # update the environments
    with open(conda_log, "wt") as logfile:

        error_code = update_mamba_conda(config, logfile, conda_base, noop)

        ec, installed = install_environments(conda_base, forge_url, logfile, install_envs, noop)

        error_code = ec or error_code

        # only update an environment if it did not install this run
        error_code = update_environments(conda_base, forge_url, logfile, update_envs - installed, noop) or error_code
        error_code = create_links(conda_base, links_list, noop) or error_code

        if remove_missing:
            remove_envs = remove_envs.union(find_unreferenced_environments(conda_base, install_envs.union(link_envs)))

        error_code = remove_environments(conda_base, remove_envs, noop) or error_code
        clean_cache(conda_base, noop)

    return error_code

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Manage CDS conda environments')

    parser.add_argument('--config-path', dest='config_path', help='path to configuration yaml file',
                        required=True)

    parser.add_argument('--error-code', '-e', dest='error_code', action='store_true',
        help='return a non-zero error code on any error, including failure to update an environment')

    parser.add_argument('--noop', action='store_true', help="Don't perform any operations on conda environments.  Just print out what would be done.")

    args = parser.parse_args()
    ec = run(args.config_path, args.noop)
    if args.error_code:
        sys.exit(ec)

