#!python

from argparse import ArgumentParser
import utilix
from admix.downloader import download, download_1t
from admix.utils import make_did

HAVE_STRAXEN = False
HAVE_CUTAX = False

try:
    import straxen
    HAVE_STRAXEN = True
except ImportError:
    print("Warning: No installation of straxen found -- you will have to pass straxen versions yourself")

try:
    import cutax
    HAVE_CUTAX = True
except ImportError:
    print("Warning: No installation of cutax found -- you will have to pass context names yourself")


def main():
    parser = ArgumentParser("admix-download")

    parser.add_argument("number", type=int, help="Run number to download")
    parser.add_argument("dtype", help="Data type to download")
    parser.add_argument("--chunks", nargs="*", help="Space-separated list of chunks to download.")
    parser.add_argument("--dir", help="Path to put the downloaded data.", default='.')
    parser.add_argument('--tries', type=int, help="Number of tries to download the data.", default=3)
    parser.add_argument('--rse', help='RSE to download from')
    parser.add_argument('--threads', help='Number of threads to use', default=3, type=int)
    parser.add_argument('--context', help='strax context you need -- this determines the hash',
                         default='xenonnt_online')
    parser.add_argument('--straxen_version', help='straxen version', default=None)
    parser.add_argument('--experiment', help="xent or xe1t", choices=['xe1t', 'xent'], default='xent')

    args = parser.parse_args()

    if args.experiment == 'xent':
        # if a particular straxen version is passed, use that
        if args.straxen_version:
            utilix_db = utilix.DB()
            hash = utilix_db.get_hash(args.context, args.dtype, args.straxen_version)
        # otherwise try to use system
        else:
            system_context = None
            # if we have cutax, try to get context from there
            err_msg = f"Could not find a useable context called {args.context} in cutax nor straxen"
            if HAVE_CUTAX:
                try:
                    system_context = getattr(cutax, args.context)()
                except:
                    try:
                        system_context = getattr(straxen, args.context)()
                    except:
                        raise RuntimeError(err_msg)
            else:
                try:
                    system_context = getattr(straxen, args.context)()
                except:
                    raise RuntimeError(err_msg)
            hash = system_context.provided_dtypes()[args.dtype]['hash']

        if args.chunks:
            chunks = [int(c) for c in args.chunks]
        else:
            chunks = None

        did = make_did(args.number, args.dtype, hash)

        downloaded = download(did, chunks=chunks, location=args.dir, tries=args.tries,
                              rse=args.rse, num_threads=args.threads)

        print(f"Download of {len(downloaded)} files finished.")

    elif args.experiment == 'xe1t':
        download_1t(args.number, args.dtype, location=args.dir, tries=args.tries, rse=args.rse,
                    num_threads=args.threads)

main()
