#!/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__)

    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.',
                           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='.')
    arguments.add_argument('--fringe-model-id', type=str,
                           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')

    allgroup = parser.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)

    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)

    if args.allFlag:
        download_models(args.fringe_model_date,
                        args.fringe_model_folder)
    else:
        download_models(args.fringe_model_date,
                        args.fringe_model_folder,
                        model_id=args.fringe_model_id)


if __name__ == '__main__':
    main()
