#!/usr/bin/env python3

"""
cli client for https://openneuro.org

(this is an alternative to https://github.com/OpenNeuroOrg/openneuro/blob/master/packages/openneuro-cli/)
"""

## TODO
#
# [ ] Factor upload()/download(); they share a lot of the same logic, just in opposite directions

import os
import io
import random
import posixpath
import logging

import json
from getpass import getpass

import click
from progressbar import *

import openneuro


class ProgressFile(io.IOBase):
    """
    A shim which reports progress in reading/writing.

    >>> fp = open("yourfile", "rb")
    >>> fp = ProgressFile(fp, lambda pos: print(f'Now at: {pos}'))
    >>> with open("yourfile (copy)", "wb") as out:
    ...    for chunk in ....

    The motivation for this is to extract progress information
    from deep inside libraries that wouldn't normally have a way
    to report it.
    """
    def __init__(self, stream, callback):
        self.wrapped = stream
        self.callback = callback

    def read(self, size, *args, **kwargs):
        try:
            return self.wrapped.read(size, *args, **kwargs)
        finally:
            self.callback(self.wrapped.tell())

    def write(self, *args, **kwargs):
        try:
            return self.wrapped.write(*args, **kwargs)
        finally:
            self.callback(self.wrapped.tell())

    def __getattr__(self, name):
        # recall: this is only run as a fallback
        return getattr(self.wrapped, name)

def credentials():
    """
    Try to load credentials to access OpenNeuro.

    Environment variables
     OPENNEURO_SERVER
     OPENNEURO_TOKEN
    has the highest precedence, followed by JSON file
     ~/.openneuro
    if either can't be found, returns None for its value.

    Returns (server, token)
    """
    server, token = 'https://openneuro.org', None

    try:
        with open(os.path.expanduser("~/.openneuro"), "r") as fp: # TODO: Windows??
            obj = json.load(fp) # XXX is this deserialization dangerous?
            server = obj.get('server', server)
            token = obj.get('apikey', token)
    except:
        #logging.exception("Couldn't load credentials from ~/.openneuro") # DEBUG
        pass

    server = os.environ.get('OPENNEURO_SERVER', server)
    token = os.environ.get('OPENNEURO_TOKEN', token)

    return server, token

# ----------------- cli ---------------

class UnsortedGroup(click.Group):
    """
    """
    # https://stackoverflow.com/questions/47972638/how-can-i-define-the-order-of-click-sub-commands-in-help
    def list_commands(self, *args, **kwargs):
        return list(self.commands)

# TODO: use click.CommandCollection <https://click.palletsprojects.com/en/7.x/commands/#merging-multi-commands>
# to avoid code duplication by making a single group whose main code does the logging in / etc and passes it along as ctx?

@click.group(cls=UnsortedGroup)
def cli():
    """
    Terminal interface to https://openneuro.org
    """
    pass

@cli.command()
@click.option("--server", "-s", default="https://openneuro.org", help="Specify server to connect to.")
def login(server):
    """
    Login to openneuro using your API token.
    Logging in is optional. It is only necessary for
    'upload' or for 'download' from an unpublished dataset.

    You can get an auth token from https://openneuro.org/keygen.

    You can override these settings with environment variables:

    \b
     OPENNEURO_SERVER - set the openneuro server instance to talk to
     OPENNEURO_TOKEN  - set the openneuro authentication token to use

    For example, you might do:

    \b
    $ OPENNEURO_TOKEN=$(pass openneuro.org-api) openneuro-cli upload ds000001 ./

    to keep your token secure in a password manager in the clear on your disk.
    """
    token = getpass("OpenNeuro API Token: ")

    # TODO: test credentials by connecting to server and trying to use them before accepting them

    with open(os.path.expanduser("~/.openneuro"), "w") as fp: # TODO: Windows??
        # this format is compatible with npm:openneuro-cli
        json.dump({'url': server, 'apikey': token}, fp)

@cli.command()
@click.option("--delete", is_flag=True, help="Delete remote files that do not exist locally.")
@click.option("-f", "--force", is_flag=True, help="Automatically create a new dataset if needed.")
@click.argument("dataset", required=False)
@click.argument("path", required=False, type=click.Path(exists=True, file_okay=False, dir_okay=True, readable=True))
def upload(delete, force, dataset, path):
    """
    Upload the contents of PATH to DATASET.

    If PATH is not given, the current directory is assumed.

    If DATASET is not given, a new dataset will be created.
    """
    if path is None:
        path = "."

    server, token = credentials()
    client = openneuro.Client(token, server=server)

    if dataset is None:
        # make a new dataset
        def prompt():
            while True: # there's gotta be a prompt-until-match in the stdlib somewhere
                r = input("Make a new dataset? [y/N] ").strip()
                if not r:
                    r = 'N'
                if r.upper() == 'Y':
                    return True
                elif r.upper() == 'N':
                    return False
                else:
                    continue # bad input, try again
        if force or prompt():
            dataset = client.createDataset()
            print(f'Created {client.server}/datasets/{dataset}')
        else:
            return # no dataset means nothing to do

    remote_files = client.files(dataset)
    remote_files = {m['filename']: m for m in remote_files}

    for dir, _, files in os.walk(path):
        for file in files:
            # NB: we have to use posixpath on the remote side, but os.path locally
            local_fname = os.path.join(dir, file)
            remote_fname = os.path.relpath(posixpath.join(dir, file), path)

            if remote_fname in remote_files:
                try:
                    if remote_files[remote_fname]['size'] == os.stat(local_fname).st_size:
                        # already uploaded
                        # (openneuro doesn't key files by hash, https://github.com/OpenNeuroOrg/openneuro/issues/805#issuecomment-648367628
                        #  so the only thing we can do is check the filesize to guess if it's the same or not :( )
                        continue
                finally:
                    del remote_files[remote_fname]

            try:
                widgets = [f'{remote_fname}: ', Percentage(), ' ', Bar(),
                           ' ', ETA(), ' ', FileTransferSpeed()]
                with open(local_fname, 'rb') as fp:
                    pbar = ProgressBar(widgets=widgets, maxval=os.stat(local_fname).st_size+1)
                     # the +1 is a workaround for 0-length files: https://github.com/niltonvolpato/python-progressbar/issues/72
                    fp = ProgressFile(fp, pbar.update)
                    pbar.start()
                    # this progress bar hack is fine for now
                    # but if we switch to it might not scale, because it requires pre-opening every single file
                    client.uploadFile(dataset, fp, remote_fname)
                    pbar.finish()
            except KeyboardInterrupt:
                raise
            except:
                logging.exception("Unable to upload {local_fname}")
                continue
    if delete:
        # go through the leftover remote_files
        # TODO: batch this with deleteFiles() ; this is *extremely slow* to do one by one; I think because internally they are doing `datalad save` (i.e. git annex add . && git add . && git commit) on every deletion.
        # TODO: compare how rsync displays this case; does it intermix the deletions among the transfers?
        for remote_fname in remote_files:
            try:
                widgets = [f'Delete {remote_fname}: ', Percentage(), ' ', Bar()]
                pbar = ProgressBar(widgets=widgets)
                pbar.start()
                client.deleteFile(dataset, remote_fname)
                pbar.finish()
            except KeyboardInterrupt:
                raise
            except:
                logging.exception("Unable to delete {remote_fname}")
                continue

@cli.command()
@click.option("--delete", is_flag=True, help="Delete local files that do not exist remotely.")
@click.argument("dataset")
@click.option("--version", "-v", help="Dataset version (e.g. '1.2.9').")
@click.argument("path", required=False, type=click.Path(exists=True, file_okay=False, dir_okay=True, writable=True))
def download(delete, version, dataset, path):
    """
    Download the contents of DATASET to PATH.

    If PATH is not given, the current directory is assumed.
    """
    if path is None:
        path = "."

    server, token = credentials()
    client = openneuro.Client(token, server=server)

    remote_files = client.files(dataset, version=version)
    remote_files = {m['filename']: m for m in remote_files}

    for remote_fname in remote_files:
        m = remote_files[remote_fname]
        local_fname = os.path.join(path, remote_fname) # XXX fix the directory traversal vuln here
            # XXX *also* this isn't windows compatible; I need to parse m['filename'] using posixpath
            # and reconstruct it using os.path

        if os.path.exists(local_fname):
            if os.stat(local_fname).st_size == m['size']:
                # already downloaded
                # (openneuro doesn't key files by hash, https://github.com/OpenNeuroOrg/openneuro/issues/805#issuecomment-648367628
                #  so the only thing we can do is check the filesize to guess if it's the same or not :( )
                continue

        os.makedirs(os.path.dirname(local_fname), exist_ok=True)
        try:
            widgets = [f'{local_fname}: ', Percentage(), ' ', Bar(),
                       ' ', ETA(), ' ', FileTransferSpeed()]
            with open(local_fname, 'wb') as fp:
                pbar = ProgressBar(widgets=widgets, maxval=m['size']+1)
                 # the +1 is a workaround for 0-length files: https://github.com/niltonvolpato/python-progressbar/issues/72
                 # TODO: sanitize the input here: make sure it's a >=0 int.
                fp = ProgressFile(fp, pbar.update)
                pbar.start()
                # ugh this is gross: because openneuro doesn't provide a good way
                # to get per-file metadata -- it expects you to batch download it
                # -- we have to call the private _download helper.
                # maybe this can be fixed? add a cache keyed by file ID?
                client._download(random.choice(m['urls']), fp)
                pbar.finish()
        except KeyboardInterrupt:
            raise
        except Exception:
            # log it but continue
            logging.exception(f"Unable to download {remote_fname}")
            continue

    if delete:
        for dir, _, files in os.walk(path):
            for file in files:
                # NB: we have to use posixpath on the remote side, but os.path locally
                local_fname = os.path.join(dir, file)
                remote_fname = os.path.relpath(posixpath.join(dir, file), path)
                if remote_fname not in remote_files:
                    try:
                        widgets = [f'Delete {local_fname}: ', Percentage(), ' ', Bar()]
                        pbar = ProgressBar(widgets=widgets)
                        pbar.start()
                        os.unlink(local_fname)
                        pbar.finish()
                    except KeyboardInterrupt:
                        raise
                    except:
                        logging.exception("Unable to delete {local_fname}")
                        continue

@cli.command()
@click.argument("dataset")
def publish(dataset):
    """
    Make DATASET publically available.

    On the main https://openneuro.org instance, this will also
    cause it to be exported to https://github.com/OpenNeuroDatasets.
    """
    server, token = credentials()
    client = openneuro.Client(token, server=server)

    client.publishDataset(dataset)

if __name__ == '__main__':
    cli()
