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

import os
import sys
import requests
import json
import getpass
import urllib3
import time
import socket
from multiprocessing import Pool
from argparse import ArgumentParser

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


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

    https://earthexplorer.usgs.gov/inventory/documentation/json-api?version=1.4.1

    """

    def __init__(self, version='1.4.1'):
        self.baseurl = 'https://earthexplorer.usgs.gov/inventory/json/v/%s/' % version

    def _api(self, endpoint='login', body=None):
        body = {'jsonRequest': json.dumps(body)} if body else {}
        r = requests.post(self.baseurl + endpoint, data=body)
        r.raise_for_status()
        dat = r.json()
        if dat.get('error'):
            sys.stderr.write(': '.join([dat.get('errorCode'), dat.get('error')]))
            exit(1)
        return dat

    @classmethod
    def login(cls, username, password=None, **kwargs):
        if password is None:
            password = getpass.getpass('Password (%s): ' % username)
        payload = {'username': username, 'password': password}
        return cls()._api('login', payload).get('data')

    @classmethod
    def search(cls, **kwargs):
        return cls()._api('search', kwargs).get('data')

    @classmethod
    def download(cls, **kwargs):
        return cls()._api('download', kwargs).get('data')

    @staticmethod
    def additionalCriteriaValues(h=None, v=None, s=None, gd=None, p=None, r=None, tile_number=None):
        k = 'additionalCriteria'
        additional = {k: {"filterType": "and", "childFilters": []}}
        if h:
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 21787, "value": h})
        if v:
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 21788, "value": v})
        if s:
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 21792, "value": s})
        if gd:
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 21795, "value": gd})
        if p:
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 21989, "value": p})
        if r:
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 19879, "value": r})
        if tile_number:  # For Sentinel-2 A,B
            additional[k]['childFilters'].append({"filterType": "value", "fieldId": 18701, "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 {"temporalFilter": {"dateField": "search_date", "startDate": sd, "endDate": ed}}


def download_url(x):
    # We need to get the redirect URL first
    fileurl, directory = x[0], x[1]
    head = requests.head(fileurl, timeout=60)
    location = head.headers.get('Location')

    base_url = fileurl.split('/download-staging')[0]
    fileurl = base_url + location
    head = requests.head(fileurl, timeout=60)

    filename = head.headers['Content-Disposition'].split('filename=')[-1].strip('"')

    local_fname = os.path.join(directory, filename)
    if os.path.exists(local_fname):
        sys.stdout.write('Already exists: %s \n' % local_fname); sys.stdout.flush()
        return

    file_size = None
    if 'Content-Length' in head.headers:
        file_size = int(head.headers['Content-Length'])
    bytes_recv = 0
    if os.path.exists(local_fname + '.part'):
        bytes_recv = os.path.getsize(local_fname + '.part')

    sys.stdout.write("Downloading %s ... \n" % local_fname); sys.stdout.flush()
    resume_header = {'Range': 'bytes=%d-' % bytes_recv}
    sock = requests.get(fileurl, headers=resume_header, timeout=60,
                        stream=True, verify=False, allow_redirects=True)

    start = time.time()
    f = open(local_fname + '.part', 'ab')
    bytes_in_mb = 1024*1024
    for block in sock.iter_content(chunk_size=bytes_in_mb):
        if block:
            f.write(block)
            bytes_recv += len(block)
    f.close()
    ns = time.time() - start
    mb = bytes_recv/float(bytes_in_mb)
    sys.stdout.write("%s (%3.2f (MB) in %3.2f (s), or  %3.2f (MB/s)) \n" % (filename, mb, ns, mb/ns)); sys.stdout.flush()

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


def download_url_wrapper(x):
    succeeded = False
    tries = 0
    while not succeeded:
        try:
            tries = tries + 1
            download_url(x)
            time.sleep(1)
            succeeded = True
        except Exception as e:
            sys.stdout.write('\n\n *** Failed download %s: %s \n' % (x, str(e))); sys.stdout.flush()
            time.sleep(60)
            if tries > 1000:
                break



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



def download_files(directory, h=None, v=None, p=None, r=None, username=None,password=None, dataset=None, tile_number=None,
                   sensor=None, N=50000, products=None, acq_date=None, gen_date=None,
                   batch=1000, threads=40, search_only=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]
            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]
            batch: How many URLs to request before working on downloads
            threads: Number of download threads to launch in parallel
            search_only: Boolean, if true only show results, don't download

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

    api_key = EarthExplorer.login(username=username,password=password)
    search = dict(apiKey=api_key, datasetName=dataset, maxResults=N)
    if any([h, v, p, r, sensor, gen_date, tile_number]):
        search.update(EarthExplorer.additionalCriteriaValues(h=h, v=v, p=p, r=r, s=sensor, gd=gen_date, tile_number=tile_number))
    if acq_date:
        search.update(EarthExplorer.temporalCriteria(ad=acq_date))
    sys.stdout.write('Full M2M Search Parameters: {}\n'.format(search))
    results = EarthExplorer.search(**search)

    n_results = results['totalHits']
    product_ids = results['results']

    sys.stdout.write('Total search results: %d \n' % n_results); sys.stdout.flush()
    if len(product_ids) < 1:
        return

    if not search_only:
        if not os.path.exists(directory):
            os.makedirs(directory,exist_ok=True)

        for pids in chunkify(product_ids, batch):


            entities = [p['entityId'] for p in pids]
            results = EarthExplorer.download(apiKey=api_key, datasetName=dataset,
                                             products=products.split(','), entityIds=entities)
            urls = [(r['url'], directory) for r in results]
            urls = [u for u in urls if all(u)]
            for url in urls:
                download_url_wrapper(url)

            #pool.map_async(download_url_wrapper, urls).get(1800)
    else:
        sys.stdout.write('RESULTS: ')
        sys.stdout.write('{}'.format(product_ids))

    return None


