import getpass
import logging
import math
import netrc
import os
import re
import subprocess
from multiprocessing.pool import ThreadPool

import numpy as np
import requests

from . import utils

try:
    input = raw_input  # Check for python 2
except NameError:
    pass

logger = logging.getLogger("sardem")


def _get_username_pass():
    """If netrc is not set up, get command line username and password"""
    print("====================================================================")
    print("Please enter NASA Earthdata credentials to download NASA hosted STRM.")
    print("See https://urs.earthdata.nasa.gov/users/new for signup info.")
    print("Or choose data_source=AWS for Mapzen tiles.")
    print("===========================================")
    username = input("Username: ")
    password = getpass.getpass(prompt="Password (will not be displayed): ")
    save_to_netrc = input(
        "Would you like to save these to ~/.netrc (machine={}) for future use (y/n)?  ".format(
            Downloader.NASAHOST
        )
    )
    return username, password, save_to_netrc.lower().startswith("y")


class Netrc(netrc.netrc):
    """Handles saving of .netrc file, fixes bug in stdlib older versions

    https://bugs.python.org/issue30806
    Uses ideas from tinynetrc
    """

    def format(self):
        """Dump the class data in the format of a .netrc file.

        Fixes issue of including single quotes for username and password"""
        rep = ""
        for host in self.hosts.keys():
            attrs = self.hosts[host]
            rep += "machine {host}\n\tlogin {attrs[0]}\n".format(host=host, attrs=attrs)
            if attrs[1]:
                rep += "\taccount {attrs[1]}\n".format(attrs=attrs)
            rep += "\tpassword {attrs[2]}\n".format(attrs=attrs)
        for macro in self.macros.keys():
            rep += "macdef {macro}\n".format(macro=macro)
            for line in self.macros[macro]:
                rep += line
            rep += "\n"

        return rep

    def __repr__(self):
        return self.format()

    def __str__(self):
        return repr(self)


class Tile:
    """class to handle tile name formation and parsing"""

    def __init__(self, left, bottom, right, top):
        self.bounds = (left, bottom, right, top)

    @staticmethod
    def get_tile_parts(tile_name):
        """Parses the lat/lon information of a .hgt tile

        Validates that the string is an actual tile name

        Args:
            tile_name

        Returns:
            tuple: lat_str (either 'N' for north, 'S' for south), lat: int latitude from 0 to 90
                lon_str (either 'W' for west, 'E' for east), lon: int longitude from 0 to 180

        Raises:
            ValueError: if regex match fails on tile_name

        Examples:
            >>> Tile.get_tile_parts('N19W156')
            ('N', 19, 'W', 156)
            >>> Tile.get_tile_parts('S5E6')
            ('S', 5, 'E', 6)
            >>> Tile.get_tile_parts('Notrealname')
            Traceback (most recent call last):
               ...
            ValueError: Invalid SRTM1 tile name: Notrealname, must match \
([NS])(\d{1,2})([EW])(\d{1,3})
        """
        lon_lat_regex = r"([NS])(\d{1,2})([EW])(\d{1,3})"
        match = re.match(lon_lat_regex, tile_name)
        if not match:
            raise ValueError(
                "Invalid SRTM1 tile name: {}, must match {}".format(
                    tile_name, lon_lat_regex
                )
            )

        lat_str, lat, lon_str, lon = match.groups()
        return lat_str, int(lat), lon_str, int(lon)

    @staticmethod
    def srtm1_tile_corner(lon, lat):
        """Integers for the bottom right corner of requested lon/lat

        Examples:
            >>> Tile.srtm1_tile_corner(3.5, 5.6)
            (3, 5)
            >>> Tile.srtm1_tile_corner(-3.5, -5.6)
            (-4, -6)
        """
        return int(math.floor(lon)), int(math.floor(lat))

    def srtm1_tile_names(self):
        """Iterator over all tiles needed to cover the requested bounds

        Args:
            None: bounds provided to Tile __init__()

        Yields:
            str: tile names that cover all of bounding box
                yielded in order of top left to bottom right

        Examples:
            >>> bounds = (-155.7, 19.1, -154.7, 19.7)
            >>> d = Tile(*bounds)
            >>> from types import GeneratorType  # To check the type
            >>> type(d.srtm1_tile_names()) == GeneratorType
            True
            >>> list(d.srtm1_tile_names())
            ['N19W156', 'N19W155']
            >>> bounds = [-156.0, 19.0, -154.0, 20.0]  # Show int bounds
            >>> list(d.srtm1_tile_names())
            ['N19W156', 'N19W155']
        """
        left, bottom, right, top = self.bounds
        left_int, top_int = self.srtm1_tile_corner(left, top)
        right_int, bot_int = self.srtm1_tile_corner(right, bottom)
        # If exact integer was requested for top/right, assume tile with that number
        # at the top/right is acceptable (dont download the one above that)
        if isinstance(top, int) or int(top) == top:
            top_int -= 1
        if isinstance(right, int) or int(right) == right:
            right_int -= 1

        # Now iterate in same order in which they'll be stithced together
        for ilat in range(top_int, bot_int - 1, -1):  # north to south
            hemi_ns = "N" if ilat >= 0 else "S"
            lat_str = "{}{:02d}".format(hemi_ns, abs(ilat))
            for ilon in range(left_int, right_int + 1):  # West to east
                hemi_ew = "E" if ilon >= 0 else "W"
                lon_str = "{}{:03d}".format(hemi_ew, abs(ilon))

                yield "{lat_str}{lon_str}".format(lat_str=lat_str, lon_str=lon_str)


class Downloader:
    """Class to download and save SRTM1 tiles to create DEMs

    Attributes:
        tile_names (iterator): strings of .hgt tiles (e.g. [N19W155.hgt])
        data_url (str): Base url where .hgt tiles are stored
        compress_type (str): format .hgt files are stored in online
        data_source (str): choices: NASA, AWS. See module docstring for explanation of sources
        cache_dir (str): explcitly specify where to store .hgt files

    Raises:
        ValueError: if data_source not a valid source string

    """

    DATA_URLS = {
        "NASA": "http://e4ftl01.cr.usgs.gov/MEASURES/SRTMGL1.003/2000.02.11",
        "NASA_WATER": "http://e4ftl01.cr.usgs.gov/MEASURES/SRTMSWBD.003/2000.02.11",
        "AWS": "https://s3.amazonaws.com/elevation-tiles-prod/skadi",
        "COP": "https://copernicus-dem-30m.s3.amazonaws.com/{t}/{t}.tif"
    }
    VALID_SOURCES = DATA_URLS.keys()
    TILE_ENDINGS = {
        "NASA": ".SRTMGL1.hgt",
        "NASA_WATER": ".SRTMSWBD.raw",
    }
    COMPRESS_TYPES = {"NASA": "zip", "NASA_WATER": "zip", "AWS": "gz"}
    NASAHOST = "urs.earthdata.nasa.gov"

    def __init__(
        self, tile_names, data_source="NASA", netrc_file="~/.netrc", cache_dir=None
    ):
        self.tile_names = tile_names
        self.data_source = data_source
        if data_source not in self.VALID_SOURCES:
            raise ValueError(
                "data_source must be one of: {}".format(",".join(self.VALID_SOURCES))
            )
        self.data_url = self.DATA_URLS[data_source]
        self.ext_type = "raw" if data_source == "NASA_WATER" else "hgt"
        self.compress_type = self.COMPRESS_TYPES[data_source]
        self.netrc_file = os.path.expanduser(netrc_file)
        self.cache_dir = cache_dir or utils.get_cache_dir()

    def _get_netrc_file(self):
        return Netrc(self.netrc_file)

    def _has_nasa_netrc(self):
        try:
            n = self._get_netrc_file()
            # Check account exists, as well is having username and password
            return (
                self.NASAHOST in n.hosts
                and n.authenticators(self.NASAHOST)[0]
                and n.authenticators(self.NASAHOST)[2]
            )
        except (OSError, IOError):
            return False

    @staticmethod
    def _nasa_netrc_entry(username, password):
        """Create a string for a NASA urs account in .netrc format"""
        outstring = "machine {}\n".format(Downloader.NASAHOST)
        outstring += "\tlogin {}\n".format(username)
        outstring += "\tpassword {}\n".format(password)
        return outstring

    def handle_credentials(self):
        """Prompt user for NASA username/password, store as attribute or .netrc

        If the user wants to save as .netrc, add to existing, or create new ~/.netrc
        """
        username, password, do_save = _get_username_pass()
        if do_save:
            try:
                # If they have a netrc existing, add to it
                n = self._get_netrc_file()
                n.hosts[self.NASAHOST] = (username, None, password)
                outstring = str(n)
            except (OSError, IOError):
                # Otherwise, make a fresh one to save
                outstring = self._nasa_netrc_entry(username, password)

            with open(self.netrc_file, "w") as f:
                f.write(outstring)
        else:
            # Save these as attritubes for the NASA url request
            self.username = username
            self.password = password

    def _form_tile_url(self, tile_name):
        """Form the url for a .hgt tile from NASA or AWS

        Args:
            tile_name (str): string name of tile
            e.g. N06W001.SRTMGL1.hgt.zip (usgs) or N19/N19W156.hgt.gz (aws)

        Returns:
            url: formatted url string with tile name

        Examples:
            >>> d = Downloader(['N19W156', 'N19W155'], data_source='NASA')
            >>> print(d._form_tile_url('N19W155'))
            http://e4ftl01.cr.usgs.gov/MEASURES/SRTMGL1.003/2000.02.11/N19W155.SRTMGL1.hgt.zip

            >>> d = Downloader(['N19W156', 'N19W155'], data_source='NASA_WATER')
            >>> print(d._form_tile_url('N19W155'))
            http://e4ftl01.cr.usgs.gov/MEASURES/SRTMSWBD.003/2000.02.11/N19W155.SRTMSWBD.raw.zip

            >>> d = Downloader(['N19W156', 'N19W155'], data_source='AWS')
            >>> print(d._form_tile_url('N19W155'))
            https://s3.amazonaws.com/elevation-tiles-prod/skadi/N19/N19W155.hgt.gz
        """
        if self.data_source == "AWS":
            lat_str, lat_int, _, _ = Tile.get_tile_parts(tile_name)
            url = "{base}/{lat}/{tile}.{ext}.{comp}".format(
                base=self.data_url,
                lat=lat_str + str(lat_int),
                tile=tile_name,
                ext=self.ext_type,
                comp=self.compress_type,
            )
        elif self.data_source.startswith("NASA"):
            url = "{base}/{tile}.{ext}".format(
                base=self.data_url,
                tile=tile_name + self.TILE_ENDINGS[self.data_source],
                ext=self.compress_type,
            )
        return url

    def _download_hgt_tile(self, url):
        """Example from https://lpdaac.usgs.gov/data_access/daac2disk "command line tips" """
        # Using AWS or a netrc file are the easy cases
        logger.info("Downloading {}".format(url))
        if self.data_source == "AWS":
            response = requests.get(url)
        elif self.data_source.startswith("NASA") and self._has_nasa_netrc():
            logger.info("Using netrc file: %s", self.netrc_file)
            response = requests.get(url)
        else:
            # NASA without a netrc file needs special auth session handling
            with requests.Session() as session:
                session.auth = (self.username, self.password)
                r1 = session.request("get", url)
                # NASA then redirects to
                # urs.earthdata.nasa.gov/oauth/authorize?scope=uid&app_type=401&client_id=...
                response = session.get(r1.url, auth=(self.username, self.password))

        return response

    def _unzip_file(self, filepath):
        """Unzips in place the .hgt files downloaded"""
        ext = os.path.splitext(filepath)[1]
        if ext == ".gz":
            unzip_cmd = ["gunzip"]
        elif ext == ".zip":
            # -o forces overwrite without prompt, -d specifices unzip directory
            unzip_cmd = "unzip -o -d {}".format(self.cache_dir).split(" ")
        subprocess.check_call(unzip_cmd + [filepath])

    def download_and_save(self, tile_name):
        """Download and save one single tile

        Args:
            tile_name (str): string name of tile
            e.g. N06W001.SRTMGL1.hgt.zip (usgs) or N19/N19W156.gz (aws)

        Returns:
            bool: True/False for Success/Failure of download
        """
        # keep all in one folder, compressed
        local_filename = self._filepath(tile_name)
        if os.path.exists(local_filename):
            logger.info("{} already exists, skipping.".format(local_filename))
        else:
            # On AWS these are gzipped: download, then unzip
            local_filename += ".{}".format(self.compress_type)
            with open(local_filename, "wb") as f:
                url = self._form_tile_url(tile_name)
                response = self._download_hgt_tile(url)
                # Now check response for auth issues/ errors
                if response.status_code == 404:
                    logger.warning("Cannot find url %s, using zeros for tile." % url)
                    # Raise only if we want to kill everything
                    # response.raise_for_status()
                    local_filename = os.path.splitext(local_filename)[0]
                    self._write_zeros(local_filename)
                    return local_filename

                f.write(response.content)
                logger.info("Writing to {}".format(local_filename))
            logger.info("Unzipping {}".format(local_filename))
            self._unzip_file(local_filename)
            # Now get rid of the .zip again
            local_filename = os.path.splitext(local_filename)[0]
        # True indicates success for this tile_name
        return local_filename

    def _filepath(self, tile_name):
        return os.path.join(self.cache_dir, tile_name + "." + self.ext_type)

    def _all_files_exist(self):
        filepaths = [self._filepath(tile_name) for tile_name in self.tile_names]
        return all(os.path.exists(f) for f in filepaths)

    def _write_zeros(self, local_filename):
        shape = (3601, 3601)
        if self.data_source == "NASA_WATER":
            dtype = np.uint8
            data = np.ones(shape, dtype=dtype) * 255
        else:
            dtype = np.int16
            data = np.zeros(shape, dtype=dtype)
        data.tofile(local_filename)

    def download_all(self):
        """Downloads and saves all tiles from tile list"""
        # Only need to get credentials for this case:
        if (
            not self._all_files_exist()
            and self.data_source.startswith("NASA")
            and not self._has_nasa_netrc()
        ):
            self.handle_credentials()

        pool = ThreadPool(processes=5)
        local_filenames = pool.map(self.download_and_save, self.tile_names)
        pool.close()
        if not any(local_filenames):
            raise ValueError(
                "No successful .hgt tiles found and downloaded:"
                " check lats/ lons of DEM box for valid SRTM land area"
                " (<60 deg latitude not open ocean)."
            )
        return local_filenames
