#!/usr/bin/env python

from multiprocessing import ProcessError
from multiprocessing.dummy import Process
import logging as log
import os
import json
import argparse
import subprocess
import tempfile
from shutil import which
from os import path
from collections import OrderedDict

SOURCE_PATH = {
    "naifKernels": "kernels",
    "pck": "kernels/pck",
    "ck":"kernels/ck",
    "spk":"kernels/spk",
    "fk": "kernels/fk",
    "iak":"kernels/iak",
    "sclk":"kernels/sclk",
    "tspk":"kernels/tspk",
    "usgs":""
}

def find_conf():
    from pathlib import Path
    local_path = Path("rclone.conf")
    # this should be installed in scripts folder, so config is one directory up in etc 
    install_path = Path(os.path.dirname(__file__)) / '..' / "etc" / "isis" / 'rclone.conf'
    repo_path = Path(os.path.dirname(__file__)) / '..' / 'config' / 'rclone.conf'
    if local_path.exists():
        return str(local_path)
    elif repo_path.exists():
        return str(repo_path)
    elif install_path.exists():
        return str(install_path)
    else:
        raise RuntimeError("No config found, try passing it in explicitly with --config <path>")


# set log level to debug
def call_subprocess(command, redirect_stdout=True, redirect_stderr=False):
    stdout = subprocess.PIPE if redirect_stdout else None
    stderr = subprocess.PIPE if redirect_stderr else None

    if not which(command[0]):
        raise ProcessLookupError(f"{command[0]} is not installed.")

    log.debug("Invoking : %s", " ".join(command))

    with subprocess.Popen(
            command,
            stdout=stdout,
            stderr=stderr) as proc:
        (out, err) = proc.communicate()

        if out:
            log.debug("Process output: ")
            log.debug(out.decode())
        if err:
            log.warning(err.decode("utf-8").replace("\\n", "\n"))

        return {
            "code": proc.returncode,
            "out": out,
            "error": err
        }


def rclone(command, config=None, extra_args=[], redirect_stdout=True, redirect_stderr=False):
    try:
        if path.isfile(config):
            log.debug(f"rclone() : using config file {config}")

            # this is probably a config file on disk so pass it through
            config_path = config
            command_with_args = ["rclone", command, f'--config={config_path}', *extra_args]
            log.debug("Invoking : %s", " ".join(command_with_args))
            return call_subprocess(command_with_args, redirect_stdout, redirect_stderr)
        else:
            log.debug(f"rclone() : using config string {config}")

            # most likely a config string
            with tempfile.NamedTemporaryFile() as f:
                config_path = path.realpath(f.name)

                log.debug(f"USING CONFIG:\n{config}")

                f.write(config.encode())
                command_with_args = ["rclone", command, f"--config={config_path}", *extra_args]
                return call_subprocess(command_with_args, redirect_stdout, redirect_stderr)
    except ProcessLookupError as not_found_e:
        log.error("Executable not found. %s", not_found_e)
        raise ProcessLookupError("rclone is not installed, please install with: 'conda install -c conda-forge rclone'")
    except Exception as generic_e:
        message = f"Error running command. Reason: {generic_e}"
        log.exception(message)
        raise Exception(message)


def get_kernel_destination_path(source_type):

    try:
        source_path = SOURCE_PATH.get(source_type)
    except KeyError as e:
        
        raise KeyError(f"kernel path not found. Source type {source_type} is invalid")

    log.debug(f"source path for {source_type} is {source_path}" )
    return source_path


def create_rclone_arguments(destination, mission_name, ntransfers=10, rclone_kwargs=[]):
    """
    Parameters
    ----------

    destination str
            path to location where files will be copied/downloaded too

    set_of_pub : set(str)
            set of kernels that will be downloaded
    """
    log.debug(f"Creating RClone command for {mission_name}")
    mission_dir_name, source_type = mission_name.replace(":", "").split("_")
    if (mission_dir_name == "legacybase"):
        # We still want things to go into base
        mission_dir_name = "base"
    log.debug(f"Mission_dir_name: {mission_dir_name}, source_type: {source_type}")

    destination = os.path.join(destination, str(mission_dir_name).replace(":",""))
    #add kernel directory path if needed
    destination = os.path.join(destination, get_kernel_destination_path(source_type))
    extra_args=[f"{mission_name}",f"{destination}", "--progress", f"--checkers={ntransfers}", f"--transfers={ntransfers}", "--track-renames", f"--log-level={log.getLevelName(log.getLogger().getEffectiveLevel())}"]

    extra_args.extend(rclone_kwargs)
    log.debug(f"Args created: {extra_args}")
    return extra_args


def main(mission, dest, cfg_path, ntransfers, kwargs):
    """
    Parameters
    ----------

    source str
            path to location where files are being copied/downloaded from

    dest : str
               path to location where files will be copied/downloaded too
    
    kwargs : list
             list of additional args to pass into rclone
          
    """
    log.info("---Script starting---")

    log.debug(f"Using config: {cfg_path}")
    result = rclone("listremotes", config=cfg_path)
    config_sources = result.get('out').decode("utf-8").split("\n")
    if config_sources == ['']:
        log.error("Remote sources came back empty. Get more info by re-running with verbose flag.")
        quit(-1)
    log.debug(f"Remote Sources: {config_sources}")

    supported_missions = OrderedDict({})
    for source in sorted(config_sources, key=lambda x: x.split("_")[-1]):
        parsed_name = source.split("_")
        # If it is a mission, it should be in the format <mission_nam>_<source_type>
        if len(parsed_name) == 2 and parsed_name[1].replace(":","") in SOURCE_PATH.keys():
            remotes_mission_name = parsed_name[0]
            supported_missions[remotes_mission_name] = supported_missions.get(remotes_mission_name, []) + [source]

    log.debug(f"Supported missions:\n {supported_missions.keys()}")
    log.debug(f"Complete Dictionary:\n {json.dumps(supported_missions, indent=2)}")

    if not mission in supported_missions.keys() and mission.upper() not in ("ALL", "LEGACYBASE"):
        raise LookupError(f"{mission} is not in the list of supported missions: {supported_missions.keys()}")

    if mission == "legacybase":
        args = create_rclone_arguments(dest, "legacybase_usgs:", ntransfers, kwargs)
        rclone(command="copy", extra_args=args, config=cfg_path, redirect_stdout=False, redirect_stderr=False)
    elif(mission.upper() == "ALL"):
        supported_missions.pop("legacybase")
        for mission, remotes in supported_missions.items():
            for remote in remotes:
                args = create_rclone_arguments(dest, remote, ntransfers, kwargs)
                rclone(command="copy", extra_args=args, config=cfg_path, redirect_stdout=False, redirect_stderr=False)
    else:
        for remote in supported_missions[mission]:
            args = create_rclone_arguments(dest, remote,  ntransfers, kwargs)
            rclone(command="copy", extra_args=args, config=cfg_path, redirect_stdout=False, redirect_stderr=False)


if __name__ == '__main__':
    usageString = "usage: downloadIsisData mission dest [-h] [-v] [-n NUM_TRANSFERS] [--config CONFIG] [rclone flags]"
    helpString = (
    '''
    Use to copy ISIS data from USGS sources as well as public end points. 

    This script combines multiple sources using the rclone command under the hood. 
    
    After positional args, accepts arbitrary rclone flags. Run "rclone help flags" to get a full list of rclone flags. 
    
    See the ISIS repo (https://github.com/USGS-Astrogeology/ISIS3) for more info. 
    
    Examples of how to run this program:
        
        # download everything for tgo
        $ downloadIsisData tgo $ISISDATA
    
        # download everything except kernels with debug output
        $ downloadIsisData all $ISISDATA --dry-run --exclude="kernels/**" -vv
        
        # download only cks and fks
        $ downloadIsisData all $ISISDATA --include="{ck/**,fk/**}" 
    ''')
    parser = argparse.ArgumentParser(description = helpString,formatter_class=argparse.RawDescriptionHelpFormatter, usage=usageString)
    parser.add_argument('mission', help='mission for files to be downloaded. Use "all" to download everything, "legacybase" to download base data for ISIS<4.1.0.')
    parser.add_argument('dest', help='the destination to download files from source')
    parser.add_argument('-v', '--verbose', action='count', default=0)
    parser.add_argument('-n', '--num-transfers', action='store', default=10)
    parser.add_argument('--config', action='store', default=find_conf())
    args, kwargs = parser.parse_known_args()
    
    log_kwargs = {
      'format': '%(asctime)s %(levelname)-8s %(message)s',
      'level': log.WARN,
      'datefmt' : '%Y-%m-%d %H:%M:%S'
    }

    if args.verbose == 0:
        log_kwargs['level'] = log.WARN
    elif args.verbose == 1:
        log_kwargs['level'] = log.INFO
    elif args.verbose >= 2:
        log_kwargs['level'] = log.DEBUG

    log.basicConfig(**log_kwargs)
    log.debug("Additional Args:", kwargs)

    main(args.mission, args.dest, os.path.expanduser(args.config), args.num_transfers, kwargs)

