#!/usr/bin/env python3

"""
fringez-download :

Download pre-generated fringe models from the NERSC web portal.
"""
import os
import sys
import argparse
import requests
import wget
from bs4 import BeautifulSoup

NERSC_url = 'https://portal.nersc.gov/project/ptf/' \
            'iband/ztf_iband_fringe_models_'


def download_models(model_date, fringe_model_dir, model_id=None):
    source_code = requests.get(NERSC_url + model_date)
    plain_text = source_code.text
    soup = BeautifulSoup(plain_text, "html.parser")

    outdir = fringe_model_dir + '/' + model_date
    if not os.path.exists(outdir):
        os.makedirs(outdir)

    if model_id:
        print('Downloading model and model lists %s to %s' % (model_id,
                                                              outdir))
    else:
        print('Downloading all models and model lists to %s' % outdir)

    for link in soup.findAll('a'):
        href = link.get('href')
        if 'model' not in href:
            continue
        if model_id and model_id not in href:
            continue
        wget.download(NERSC_url + model_date + '/' + href, out=outdir)


def main():
    """Download pre-generated fringe models from the NERSC web portal."""

    # Get arguments
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    arguments = parser.add_argument_group('arguments')
    arguments.add_argument('--fringe-model-date', type=str,
                           help='Model date for downloading pre-generated '
                                'models from the NERSC web portal. See '
                                'fringez github page for list of available '
                                'model dates.', default='20200723',
                           required=True)
    arguments.add_argument('--fringe-model-folder', type=str,
                           help='Models are downloaded by default into the '
                                'current directory. --fringe-model-folder will '
                                'direct the models to be downloaded to this '
                                'folder instead.', default='.')

    allorone = parser.add_argument_group('all or single')
    allgroup = allorone.add_mutually_exclusive_group()
    allgroup.add_argument('--all', dest='allFlag',
                          action='store_true',
                          help='Download all fringe models in the '
                               '--fringe-model-date directory. DEFAULT.')
    allgroup.add_argument('--single', dest='allFlag',
                          action='store_false',
                          help='Download only a single fringe model from the '
                               '--fringe-model-date directory. Requires '
                               'setting the --fringe-model-id argument.')
    parser.set_defaults(allFlag=True)

    alliamges = parser.add_argument_group('arguments for --single')
    alliamges.add_argument('--fringe-model-id', type=str, default=None,
                           help='If --single is selected, only the fringe '
                                'model containing the --fringe-model-id is'
                                ' downloaded from the NERSC web portal. The '
                                '--fringe-model-id is the cid and the qid of '
                                'the fringe model combined in the following '
                                'syntax: c01_q1')

    args = parser.parse_args()

    if not args.allFlag and not args.fringe_model_id:
        print('--fringe-model-id must be set if --single is selected')
        sys.exit(0)

    download_models(args.fringe_model_date,
                    args.fringe_model_folder,
                    model_id=args.fringe_model_id)


if __name__ == '__main__':
    main()
