"""
Search and download M2M data, skip those already downloaded
"""

import os
import sys
import re
import requests
import json
import getpass
import urllib3
import time
import datetime
import logging
import traceback
import multiprocessing
from argparse import ArgumentParser

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
stream = logging.StreamHandler()
stream.setLevel(logging.INFO)
FORMAT = '%(asctime)-15s ' \
         '%(name)-15s ' \
         '%(levelname)-8s ' \
         '%(message)s '
formatter = logging.Formatter(FORMAT)
stream.setFormatter(formatter)
logger.addHandler(stream)


class M2MError(Exception):
    pass


class DDSError(Exception):
    pass


class M2M(object):
    """
    Web-Service interface for EarthExplorer JSON Machine-to-Machine API

    https://m2m.cr.usgs.gov/api/docs/json/

    """

    def __init__(self, instance='ops'):
        url_lookup = dict(
            ops='https://m2m.cr.usgs.gov/api/api/json/stable/',
            devmast='https://m2mdevmast.cr.usgs.gov/api/api/json/stable/',
            devsys='https://m2mdev.cr.usgs.gov/devsys/api/api/json/stable/'
        )
        self.baseurl = url_lookup.get(instance)
        self.product_name_dict = dict(landsat='Level-1 GeoTIFF Data Product',
                                      landsat_tm_c2_l1='Landsat Collection 2 Level-1 Product Bundle',
                                      landsat_etm_c2_l1='Landsat Collection 2 Level-1 Product Bundle',
                                      landsat_ot_c2_l1='Landsat Collection 2 Level-1 Product Bundle',
                                      landsat_tm_c2_l2='Landsat Collection 2 Level-2 Product Bundle',
                                      landsat_etm_c2_l2='Landsat Collection 2 Level-2 Product Bundle',
                                      landsat_ot_c2_l2='Landsat Collection 2 Level-2 Product Bundle',
                                      ard_tile='Surface Reflectance',
                                      sentinel='L1C Tile in JPEG2000 format'
                                      )
        self.ard_product_lookup = dict(
            SR='Surface Reflectance',
            QA='Quality Assessment',
            TOA='Top of Atmosphere',
            ST='Provisional Surface Temperature',
            BT='Brightness Temperature'
        )

    def accumulate_errors(self, response: requests.models.Response) -> dict:
        """
        Determine if errors are present in the response.  Assign an arbitrary
        level to each type of error.  The highest level error will take
        precedence when raising an exception.  All errors are logged.
        """

        errors = dict()
        server_exception = "M2M Server Exception"
        api_exception = "M2M API Exception"

        code = response.status_code
        reason = response.reason

        # InventoryConnectionException
        if code >= 500:
            msg = "{0}: \n\t{1}: \n\t{2}".format(server_exception, code, reason)
            logger.error(msg)
            errors['Level-1'] = msg

        elif code >= 400:
            msg = "{0}: \n\t{1}: \n\t{2}".format(api_exception, code, reason)
            logger.error(msg)
            errors['Level-2'] = msg

        try:
            data = response.json()

        # LTAError
        except ValueError as e:
            msg = "{0}:\n\tUnable to parse JSON response.\n\t{1}:\n\t{2}".format(api_exception, e,
                                                                                 traceback.format_exc())
            logger.error(msg)
            errors['Level-3'] = msg
            data = {'data': None}

        # LTAError
        err_code = data.get('errorCode', None)
        err_msg = data.get('errorMessage', None)

        if err_code or err_msg:
            msg = "{0}: \n\t{1}: {2}".format(api_exception, err_code, err_msg)
            logger.error(msg)
            errors['Level-4'] = msg

        # LTAError
        if 'data' not in data.keys():
            msg = "{0}: \n\tNo data found in response".format(api_exception)
            logger.error(msg)
            errors['Level-5'] = msg

        return errors

    def raise_error(self, errors: dict) -> None:
        """
        Raise the appropriate exception with message if an error was found
        in the response
        """
        levels = range(5, 0, -1)
        for level in levels:
            msg = errors.get('Level-{}'.format(level), None)
            if msg:
                if level > 1:
                    raise M2MError(msg)

    def _parse(self, response):
        """
        Attempt to parse the JSON response, which always contains additional
        information that we might not always want to look at (except on error)

        :param response: requests.models.Response
        :return: dict
        """
        self.raise_error(self.accumulate_errors(response))
        data = response.json()
        logger.debug('RESPONSE {}'.format(data))
        return data

    def _api_request(self, verb, resource, data, headers=None):
        url = self.baseurl + resource
        data_ = {k: v if k != 'password' else 'xxxxx' for k, v in data.items()}
        logger.debug('POST {} {}'.format(url, data_))
        response = getattr(requests, verb)(url, headers=headers, data=json.dumps(data))
        return self._parse(response)

    def login(self, username, password=None, **kwargs):
        if password is None:
            password = getpass.getpass('Password (%s): ' % username)
        payload = {'username': username, 'password': password}
        return self._api_request('post', 'login', payload).get('data')

    def scene_search(self, headers, data):
        return self._api_request('post', 'scene-search', data, headers).get('data')

    def get_product_names(self, dataset, products):
        lookup = self.product_name_dict
        if dataset.lower() == 'ard_tile':
            if products:
                products = products.split(',')
                products = [self.ard_product_lookup.get(p) for p in products]
            else:
                products = ['Surface Reflectance']
        else:
            products = [lookup.get(list(filter(lambda x: dataset.lower().startswith(x), lookup.keys()))[0])]

        return products

    def download_options(self, headers, entity_ids, dataset, products):
        params = {
            'entityIds': entity_ids,
            'datasetName': dataset
        }
        resp = self._api_request('post', 'download-options', params, headers)
        product_names = self.get_product_names(dataset, products)
        product_info = filter(lambda i: i['productName'].strip() in product_names, resp.get('data'))
        return [(p.get('entityId'), p.get('id')) for p in product_info]

    def download_retrieve(self, headers, entity_ids, label):
        """
        docs:
        https://m2m.cr.usgs.gov/api/docs/reference/#download-retrieve

        Make a request to the M2M download-retrieve endpoint.  This will return all
        available and previously requests, but not completed downloads.  ESPA
        must utilize this endpoint to retrieve downloads for data that required
        staging prior to being made available.  Otherwise, repeat attempts to download-request
        for the same scene will continuously return the URL as 'preparing' with the 'data-staging'
        address.

        Args:
            headers
            entity_ids
            label
        Returns:
            List[Tuple[urls, entityIds]]
        """
        params = {'downloadApplication': 'EE', 'label': label}
        resp = self._api_request('post', 'download-retrieve', params, headers)
        data = resp.get('data')
        avail = data.get('available')
        req = data.get('requested')

        avail = list(filter(lambda x: x.get('entityId') in entity_ids, avail))
        req = list(filter(lambda y: y.get('entityId') in entity_ids and y.get('statusCode') == 'D', req))

        results = list()
        results.extend(avail)
        results.extend(req)

        msg = 'Number of downloads ready: {}'.format(len(results))
        logger.info(msg)

        ready = [e.get('entityId') for e in results]
        not_ready = [e for e in entity_ids if e not in ready]
        if not_ready:
            msg = 'Still waiting on staging for {} IDs'.format(len(not_ready))
            logger.warning(msg)

        return [(r.get('url'), r.get('entityId')) for r in results]

    def download_request(self, headers, entity_ids, download_ids, label):
        params = {
            'downloads': [
                {'entityId': e,
                 'productId': i}
                for e, i in zip(entity_ids, download_ids)
            ],
            'downloadApplication': 'EE',
            'label': label
        }
        resp = self._api_request('post', 'download-request', params, headers)

        data = resp.get('data')
        avail = data.get('availableDownloads')
        prep = data.get('preparingDownloads')

        urls = list()

        if avail:
            msg = 'Number of available downloads to try: {}'.format(len(avail))
            logger.warning(msg)
            urls.extend(avail)

        if prep:
            msg = 'Number of downloads still staging: {}'.format(len(prep))
            logger.warning(msg)
            # These are download-staging URLs, we skip these because they will
            # timeout if we try to request them.
            # urls.extend(prep)
            logger.warning('Waiting 10 seconds...')
            time.sleep(10)

        urls = filter(lambda x: x, urls)

        return [x.get('url') for x in urls]

    @staticmethod
    def additionalCriteriaValues(h=None, v=None, p=None, r=None, s=None, sc=None, rg=None, gd=None, tile_number=None):
        k = 'metadataFilter'
        additional = {k: {"filterType": "and", "childFilters": []}}
        if h:
            additional[k]['childFilters'].append(
                {"filterType": "between", "filterId": "5e83a38b7e5e872c", "firstValue": h, "secondValue": h})
        if v:
            additional[k]['childFilters'].append(
                {"filterType": "between", "filterId": "5e83a38b5e3a58ab", "firstValue": v, "secondValue": v})
        if s:
            additional[k]['childFilters'].append({"filterType": "value", "filterId": "5e83a38b9bdb525f", "value": s})
        if sc:
            additional[k]['childFilters'].append({"filterType": "value", "filterId": "5e83a38b7b117c3d", "value": sc})
        if gd:
            additional[k]['childFilters'].append(
                {"filterType": "between", "filterId": "5e83a38b4357a01b", "firstValue": gd, "secondValue": gd})
        if rg:
            additional[k]['childFilters'].append({"filterType": "value", "filterId": "5e83a38b23954b26", "value": rg})
        if p:
            additional[k]['childFilters'].append(
                {"filterType": "between", "filterId": "5e83d0b81d20cee8", "firstValue": p, "secondValue": p})
        if r:
            additional[k]['childFilters'].append(
                {"filterType": "between", "filterId": "5e83d0b849ed5ee7", "firstValue": r, "secondValue": r})
        if tile_number:  # For Sentinel-2 A,B
            additional[k]['childFilters'].append(
                {"filterType": "value", "filterId": "5e83a42cc36e732d", "value": tile_number,
                 "operand": "like"})
        return additional

    @staticmethod
    def temporalCriteria(ad):
        dates = ad.split(',')
        sd, ed = dates if len(dates) == 2 else dates * 2
        return {"acquisitionFilter": {"start": sd, "end": ed}}


def handle_redirect(x):
    head = requests.head(x, timeout=60)
    location = head.headers.get('Location')
    if location is None:
        msg = 'Error retrieving redirect URL for {}'.format(x)
        logger.error(msg)
        raise DDSError(msg)

    base_url = x.split('/download-staging')[0]
    redirect_url = base_url + location
    return redirect_url


def download_url(x):
    # We need to get the redirect URL first
    pattern = r'(L[E|T|O|C]\d{2}_L\d{1}\w{2}_\d{6}_\d{8}_\d{8}_\d{2}_\w\d(\.tar)?)|(L\w{1}\d{2}_\w{2}_\d{6}_\d{8}_\d{8}_C\d{2}_V\d{2}_.*.tar)|(L1C_.*\.zip)|(S2A_.*\.zip)'
    fileurl, directory = x[0], x[1]
    bytes_recv = 0
    bytes_in_mb = 1024 * 1024

    if 'download-staging' in fileurl:
        fileurl_ = handle_redirect(fileurl)
    else:
        fileurl_ = fileurl

    head = requests.head(fileurl_, timeout=60)

    content_disposition = head.headers.get('Content-Disposition')
    location = head.headers.get('Location')

    try:
        if content_disposition:
            filename = content_disposition.split('filename=')[-1].strip('"').strip("'")
        elif location:
            filename = re.search(pattern, location).group()
        else:
            raise DDSError
    except Exception:
        msg = 'Unable to determine filename from headers (content-disposition and/or location) in {}'.format(fileurl_)
        logger.error(msg)
        raise DDSError(msg)

    # TODO can do this better, match extensions with appropriate datasets but seems only to happen for Landsat C2
    if os.path.splitext(filename)[-1] == '':
        filename = filename + '.tar'

    local_fname = os.path.join(directory, filename)
    if os.path.exists(local_fname):
        logger.warning('Already exists - skipping: %s \n' % local_fname)
        return

    #MOD5
    if "CU" not in local_fname:
        logger.warning('only downloading CU not this - skipping: %s \n' % local_fname)
        return
    
    try:
        file_size = None
        if 'Content-Length' in head.headers:
            file_size = int(head.headers['Content-Length'])

            if os.path.exists(local_fname + '.part'):
                bytes_recv = os.path.getsize(local_fname + '.part')

            logger.info("Downloading\n\t{}\n\tto\n\t{} ... \n".format(fileurl_, local_fname))
            resume_header = {'Range': 'bytes=%d-' % bytes_recv}
            sock = requests.get(fileurl_, headers=resume_header, timeout=10,
                                stream=True, verify=False, allow_redirects=True)

            start = time.time()
            with open(local_fname + '.part', 'ab') as f:

                for block in sock.iter_content(chunk_size=bytes_in_mb):
                    if block:
                        f.write(block)
                        bytes_recv += len(block)
            ns = time.time() - start
            mb = bytes_recv / float(bytes_in_mb)
            logger.info("Complete - %s (%3.2f (MB) in %3.2f (s), or  %3.2f (MB/s)) \n" % (filename, mb, ns, mb / ns))

            if bytes_recv >= file_size:
                os.rename(local_fname + '.part', local_fname)

        else:  # content-length unknown - cloud-hosted?
            if os.path.exists(local_fname + '.part'):
                # We aren't gonna deal with this if we don't know
                # the final file size
                os.remove(local_fname + '.part')
            logger.info("Downloading\n\t{}\n\tto\n\t{} ... \n".format(fileurl_, local_fname))
            sock = requests.get(fileurl_, timeout=10, stream=True, verify=False, allow_redirects=True)
            start = time.time()
            with open(local_fname, 'ab') as f:
                for block in sock.iter_content(chunk_size=bytes_in_mb):
                    if block:
                        f.write(block)
                        bytes_recv += len(block)
            ns = time.time() - start
            mb = bytes_recv / float(bytes_in_mb)
            logger.info("Complete - %s (%3.2f (MB) in %3.2f (s), or  %3.2f (MB/s)) \n" % (filename, mb, ns, mb / ns))
    except Exception as e:
        try:
            # Need to remove the *.tar/*.zip/*.tar.gz
            # if we aren't writing to *.part
            os.remove(local_fname)
        except FileNotFoundError:
            pass
        raise


retry_downloads_at_the_end = []

def download_url_wrapper(x):
    #MOD4
    succeeded = False
    tries = 0
    while not succeeded:
        try:
            tries = tries + 1
            download_url(x)
            #time.sleep(1)
            succeeded = True
            sys.stdout.write('\n\n *** Succeeded download %s: %s \n' % (x, x,)); sys.stdout.flush()
        except Exception as e:
            sys.stdout.write('\n\n *** Failed download %s: %s \n' % (x, str(e))); sys.stdout.flush()
            time.sleep(10)
            if tries > 100:
                sys.stdout.write('\n\n *** Really failed download this file may be missing now, will retry one more time at the end of script %s: %s \n' % (x, str(e))); sys.stdout.flush()
                retry_downloads_at_the_end.append(x)
                break


def chunkify(iterable, n=1):
    l = len(iterable)
    for ndx in range(0, l, n):
        yield iterable[ndx:min(ndx + n, l)]


def get_product_ids(results):
    return [x['displayId'] for x in results], [y['entityId'] for y in results]


def clean(param):
    if isinstance(param, str):
        return param.strip('"').strip("'")
    else:
        return param


def download_files(directory, h=None, v=None, p=None, r=None, username=None,password=None, dataset=None, tile_number=None,
                   sensor=None, spacecraft=None, region=None, N=50000, products=None, acq_date=None, gen_date=None,
                   threads=40, search_only=False, instance='ops', debug=False):
    """
    Search for and download files to local directory

        Args:
            directory: Relative path to local directory (will be created)
            h: Tile Grid Horizontal [Optional]
            v: Tile Grid Vertical [Optional]
            p: The WRS2 path [Optional]
            r: The WRS2 row [Optional]
            username: ERS Username (with full M2M download access) [Optional]
            dataset: EarthExplorer Catalog datasetName [e.g. ARD_TILE or SENTINAL_2A]
            tile_number: For SENTINEL-2
            sensor: Satellite instrument [All, OLI_TIRS, ETM, TM]
            spacecraft: Satellite platform [Landsat 4, 5, 7, 8]
            region: ARD Tile Grid Region [AK, CU, HI]
            N: Maximum number of search results to return
            products: Comma-delmited list of download products as a single string [e.g 'TOA,BT,SR,QA']
            acq_date: Search Date image acquired [Format: %Y-%m-%d]
            gen_date: Tile production date [Format: %Y-%m-%d]
            threads: Number of download threads to launch in parallel
            search_only: Boolean, if true only show results, don't download
            instance: What instance of M2M to use (devsys, devmast, ops)
            debug (bool): If True, set log level to DEBUG

    """
    if debug:
        stream.setLevel(logging.DEBUG)

    if not directory and not search_only:
        logger.error('Must specify download directory')
        sys.exit(0)

    if directory and not os.path.exists(directory):
        os.makedirs(directory)

    m2m = M2M(instance)

    #MOD6
    #token = m2m.login(username=username or input('Enter ERS username: '))
    token = m2m.login(username=username,password=password)
    header = {'X-Auth-Token': token}

    scene_filter = dict(
        ingestFilter=None,
        spatialFilter=None,
        metadataFilter=None,
        cloudCoverFilter={'max': 100, 'min': 0},
        acquisitionFilter=None
    )

    # if any([h, v, p, r, sensor, spacecraft, region, gen_date, tile_number]):
    #     # scene_filter.update(
    #     #     m2m.additionalCriteriaValues(h=h, v=v, p=p, r=r, s=sensor, sc=spacecraft, rg=region, gd=clean(gen_date),
    #     #                                  tile_number=tile_number))
    #     scene_filter.update(
    #         m2m.additionalCriteriaValues(h=None, v=None, p=None, r=None, s=None, sc=None, rg=None, gd=None,
    #                                      tile_number=None))

    # if acq_date:
    #     scene_filter.update(m2m.temporalCriteria(ad=clean(acq_date)))

    search = dict(
        datasetName=dataset,
        maxResults=N,
        sceneFilter=scene_filter)

    results = m2m.scene_search(header, search)

    n_results = results.get('totalHits')
    display_ids, entity_ids = get_product_ids(results.get('results'))

    logger.info('Total search results: %d \n' % n_results)

    if len(display_ids) < 1:
        logger.warning('No results found!')
        #MOD3
        #sys.exit(0)
        return

    if not search_only:
        download_info = m2m.download_options(header, entity_ids, dataset, products)
        e_ids = [x[0] for x in download_info]
        d_ids = [x[1] for x in download_info]

        urls = list()

        # Make the initial download request, then sleep for 10 seconds
        # to allow data to stage
        # List[url]
        req_avail = m2m.download_request(header, e_ids, d_ids, 'm2m_download_{}'.format(dataset))
        if req_avail:
            urls.extend([(x, directory) for x in req_avail])

        # Tuple[url, entityId]
        retr_avail = m2m.download_retrieve(header, e_ids, 'm2m_download_{}'.format(dataset))
        logger.debug('Retrieved URLs and entityIDs: {}'.format(retr_avail))
        urls.extend([(z[0], directory) for z in retr_avail])

        #MOD2
        pool = multiprocessing.Pool(threads)
        pool.map_async(download_url_wrapper, urls).get(600)
        pool.close()
        pool.join()

        #MOD
        #these are the ones that m2m reports as staged but they arent in reality
        #so they fail when we try them the first time
        retry_downloads_at_the_end_copy = retry_downloads_at_the_end.copy()
        for url in retry_downloads_at_the_end_copy:
            download_url_wrapper(url)

    else:
        now = datetime.datetime.now()
        name = datetime.datetime.strftime(now, '%Y-%m-%d_%H:%M:%S')
        name = 'results_{}.txt'.format(name)
        text_output = os.path.join(directory, name)
        with open(text_output, 'w') as f:
            for i in display_ids:
                f.write('{}\n'.format(i))
        logger.info('SEARCH RESULTS: ')
        logger.info('{}'.format(display_ids))
        logger.info('Writing results to {}'.format(text_output))

    return None

"""
#MOD1
def build_command_line_arguments():
    description = __doc__
    parser = ArgumentParser(description=description, add_help=False)
    req_parser = parser.add_argument_group('required arguments')
    parser.add_argument('--help', action='help', help='show this help message and exit')
    req_parser.add_argument('-d', '--directory', type=str, dest='directory',
                            help='Relative path to download all data')
    req_parser.add_argument('-u', '--username', type=str, dest='username', default=None,
                            help='ERS Username (with full M2M download access)')
    parser.add_argument('-t', '--threads', type=int, dest='threads', default=40,
                        help='Number of parallel download threads [Default: 40]')
    parser.add_argument('-m', '--max', type=int, dest='N', default=50000,
                        help='Maximum number of Tile results to return [Default: 50000]')
    parser.add_argument('-p', '--path', type=int, dest='p', default=None,
                        help='WRS-2 Path [Default: None]')
    parser.add_argument('-r', '--row', type=int, dest='r', default=None,
                        help='WRS-2 Row [Default: None]')
    parser.add_argument('-h', '--horizontal', type=int, dest='h', default=None,
                        help='ARD Tile Grid Horizontal [Default: All]')
    parser.add_argument('-v', '--vertical', type=int, dest='v', default=None,
                        help='ARD Tile Grid Vertical [Default: All]')
    parser.add_argument('--region', type=str, dest='region', default=None,
                        choices=['CU', 'AK', 'HI'],
                        help='ARD Tile Grid Region [Default: None]')
    parser.add_argument('--tile-number', type=str, dest='tile_number', default=None,
                        help='S2 Tile Number e.g. T19TDK [Default: None]')
    parser.add_argument('-s', '--sensor', type=str, dest='sensor', default=None,
                        choices=['All', 'OLI_TIRS', 'ETM', 'TM'],
                        help='Landsat sensor Identifier [Default: All]')
    parser.add_argument('--spacecraft', dest='spacecraft', default=None,
                        choices=[f'LANDSAT_{x}' for x in (4, 5, 7, 8)],
                        help='Landsat spacecraft identifier [Default: None]')
    parser.add_argument('--products', dest='products', default=None,
                        help='M2M ARD product names (e.g. "SR,TOA,BT")')
    req_parser.add_argument('--dataset', type=str, dest='dataset', required=True,
                            help='EE Catalog dataset [e.g. SENTINEL_2A]')
    parser.add_argument('--acq-date', type=str, dest='acq_date', default=None,
                        help='Search Date Acquired (YYYY-MM-DD)')
    parser.add_argument('--gen-date', type=str, dest='gen_date', default=None,
                        help='ARD Tile Production Date (YYYY/MM/DD)')
    parser.add_argument('--search-only', action='store_true', dest='search_only',
                        help='Only return search results, do not download')
    parser.add_argument('--instance', type=str, dest='instance', choices=['devsys', 'devmast', 'ops'], default='ops',
                        help='Which instance of M2M to use (devsys, devmast, ops)')
    parser.add_argument('--debug', action='store_true', default=False,
                        help='Set log level to DEBUG (Hint: Shows M2M POST requests and responses)')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    download_files(**vars(build_command_line_arguments()))
"""