import random
import threading
import typing
from time import time

import backoff
import ee
import requests
from te_schemas import results
from te_schemas.results import Raster
from te_schemas.results import TiledRaster
from te_schemas.results import Url
from te_schemas.schemas import BandInfoSchema
from te_schemas.schemas import CloudResults
from te_schemas.schemas import CloudResultsSchema

from . import GEEImageError
from . import GEETaskFailure

# Google cloud storage bucket for output
BUCKET = "ldmt"

# Number of minutes a GEE task is allowed to run before timing out and being
# cancelled
TASK_TIMEOUT_MINUTES = 48 * 60


def get_region(geom):
    """Return ee.Geometry from supplied GeoJSON object."""
    poly = get_coords(geom)
    ptype = get_type(geom)

    if ptype.lower() == 'multipolygon':
        region = ee.Geometry.MultiPolygon(poly)
    else:
        region = ee.Geometry.Polygon(poly)

    return region


def get_coords(geojson):
    """."""

    if geojson.get('features') is not None:
        return geojson.get('features')[0].get('geometry').get('coordinates')
    elif geojson.get('geometry') is not None:
        return geojson.get('geometry').get('coordinates')
    else:
        return geojson.get('coordinates')


def get_type(geojson):
    """."""

    if geojson.get('features') is not None:
        return geojson.get('features')[0].get('geometry').get('type')
    elif geojson.get('geometry') is not None:
        return geojson.get('geometry').get('type')
    else:
        return geojson.get('type')


class gee_task(threading.Thread):
    """Run earth engine task against the trends.earth API"""
    def __init__(self, task, prefix, logger, metadata=None):
        threading.Thread.__init__(self)
        self.task = task
        self.prefix = prefix
        self.logger = logger
        # self.metadata is used only to facilitate saving the final JSON output
        # for Trends.Earth
        self.metadata = metadata
        self.state = self.task.status().get('state')
        self.start()

    def cancel_hdlr(self, details):
        self.logger.debug(
            "GEE task {} timed out after {} hours".format(
                self.task.status().get('id'),
                (time() - self.start_time) / (60 * 60)
            )
        )
        ee.data.cancelTask(self.task.status().get('id'))

    def on_backoff_hdlr(self, details):
        details.update({'task_id': self.task.status().get('id')})
        self.logger.debug(
            "Backing off {wait:0.1f} seconds after {tries} tries "
            "calling function {target} for task {task_id}".format(**details)
        )

    def poll_for_completion(self):
        @backoff.on_predicate(
            backoff.expo,
            lambda x: x in ['READY', 'RUNNING'],
            on_backoff=self.on_backoff_hdlr,
            on_giveup=self.cancel_hdlr,
            max_time=TASK_TIMEOUT_MINUTES * 60,
            factor=3,
            base=1.4,
            max_value=600
        )
        def get_status(self):
            self.logger.send_progress(self.task.status().get('progress', 0.0))
            self.state = self.task.status().get('state')

            return self.state

        return get_status(self)

    def run(self):
        self.task.start()
        self.start_time = time()
        self.logger.debug(
            "Starting GEE task {}.".format(self.task.status().get('id'))
        )
        self.poll_for_completion()

        if not self.state:
            raise GEETaskFailure(self.task)

        if self.state == 'COMPLETED':
            self.logger.debug(
                "GEE task {} completed.".format(self.task.status().get('id'))
            )
        elif self.state == 'FAILED':
            self.logger.debug(
                "GEE task {} failed: {}".format(
                    self.task.status().get('id'),
                    self.task.status().get('error_message')
                )
            )
            raise GEETaskFailure(self.task)
        else:
            self.logger.debug(
                "GEE task {} returned status {}: {}".format(
                    self.task.status().get('id'), self.state,
                    self.task.status().get('error_message')
                )
            )
            raise GEETaskFailure(self.task)

    def get_urls(self):
        @backoff.on_exception(
            backoff.expo,
            requests.exceptions.RequestException,
            on_backoff=self.on_backoff_hdlr,
            max_time=60,
            factor=3,
            base=1.4,
            max_value=600
        )
        def request_urls(self):
            return requests.get(
                f'https://www.googleapis.com/storage/v1/b/{BUCKET}/o?prefix={self.prefix}'
            )

        resp = request_urls(self)

        if not resp or resp.status_code != 200:
            self.logger.debug(
                f'Failed to list urls for results from {self.task}'
            )
            raise GEETaskFailure(self.task)

        items = resp.json()['items']

        if len(items) < 1:
            self.logger.debug('No urls were found for {}'.format(self.task))
            raise GEETaskFailure(self.task)
        else:
            urls = []

            for item in items:
                urls.append(Url(item['mediaLink'], item['md5Hash']))

            return urls

    def get_uris(self):
        @backoff.on_exception(
            backoff.expo,
            requests.exceptions.RequestException,
            on_backoff=self.on_backoff_hdlr,
            max_time=60,
            factor=3,
            base=1.4,
            max_value=600
        )
        def request_uris(self):
            return requests.get(
                f'https://www.googleapis.com/storage/v1/b/{BUCKET}/o?prefix={self.prefix}'
            )

        resp = request_uris(self)

        if not resp or not resp.json().get('items'):
            self.logger.debug(
                f'Failed to list uris for results from {self.task}'
            )
            raise GEETaskFailure(self.task)

        items = resp.json()['items']

        if len(items) < 1:
            self.logger.debug('No uris were found for {}'.format(self.task))
            raise GEETaskFailure(self.task)
        else:
            uris = []

            for item in items:
                uris.append(
                    results.URI(
                        type='cloud',
                        uri=item['mediaLink'],
                        etag=results.Etag(
                            hash=item['md5Hash'],
                            type=results.EtagType.GCS_CRC32C
                        )
                    )
                )

            return uris


# Not using dataclass as not in python 3.6
class TEImage(object):
    "A class to store GEE images and band info for export to cloud storage"

    def __init__(self, image, band_info):
        self.image = image
        self.band_info = band_info

        self._check_validity()

    def _check_validity(self):
        if len(self.band_info) != len(self.image.getInfo()['bands']):
            raise GEEImageError(
                f'Band info length ({len(self.band_info)}) does not match '
                f'number of bands in image ({self.image.getInfo()["bands"]})'
            )

    def merge(self, other):
        "Merge with another TEImage object"
        self.image = self.image.addBands(other.image)
        self.band_info.extend(other.band_info)

        self._check_validity()

    def addBands(self, bands, band_info):
        "Add new bands to the image"
        self.image = self.image.addBands(bands)
        self.band_info.extend(band_info)

        self._check_validity()

    def selectBands(self, band_names):
        "Select certain bands from the image, dropping all others"
        band_indices = [
            i for i, bi in enumerate(self.band_info) if bi.name in band_names
        ]

        if len(band_indices) < 1:
            raise GEEImageError('Bands "{}" not in image'.format(band_names))

        self.band_info = [self.band_info[i] for i in band_indices]
        self.image = self.image.select(band_indices)

        self._check_validity()

    def setAddToMap(self, band_names=[]):
        "Set the layers that will be added to the map by default"

        for i in range(len(self.band_info)):
            if self.band_info[i].name in band_names:
                self.band_info[i].add_to_map = True
            else:
                self.band_info[i].add_to_map = False

    def export(
        self, geojsons, task_name, crs, logger, execution_id=None, proj=None
    ):
        "Export layers to cloud storage"

        if not execution_id:
            execution_id = str(random.randint(1000000, 99999999))
        else:
            execution_id = execution_id

        if not proj:
            proj = self.image.projection()

        tasks = []
        n = 1

        for geojson in geojsons:
            if task_name:
                out_name = '{}_{}_{}'.format(execution_id, task_name, n)
            else:
                out_name = '{}_{}'.format(execution_id, n)

            export = {
                'image': self.image,
                'description': out_name,
                'fileNamePrefix': out_name,
                'bucket': BUCKET,
                'maxPixels': 1e13,
                'crs': crs,
                'scale': ee.Number(proj.nominalScale()).getInfo(),
                'region': get_coords(geojson),
                'formatOptions': {
                    'cloudOptimized': True
                }
            }
            t = gee_task(
                task=ee.batch.Export.image.toCloudStorage(**export),
                prefix=out_name,
                logger=logger
            )
            tasks.append(t)
            n += 1

        logger.debug("Exporting to cloud storage.")
        urls = []

        for task in tasks:
            task.join()
            urls.extend(task.get_urls())

        gee_results = CloudResults(task_name, self.band_info, urls)
        results_schema = CloudResultsSchema()
        json_results = results_schema.dump(gee_results)

        return json_results


class GEEImage():
    def __init__(
        self,
        ee_image: ee.Image,
        bands: typing.List[results.Band],
        datatype: results.DataType = results.DataType.INT16
    ):
        self.ee_image = ee_image
        self.bands = bands
        self.datatype = datatype

        self._check_validity()

    def _check_validity(self):
        if len(self.bands) != len(self.ee_image.getInfo()['bands']):
            raise GEEImageError(
                f'Band info length ({len(self.bands)}) '
                'does not match number of bands in image '
                f'({self.ee_image.getInfo()["bands"]})'
            )

    def merge(self, other):
        "Merge with another GEEImage object"

        if self.datatype != other.datatype:
            raise GEEImageError(
                f'Attempted to merge {self.datatype} image with '
                f'{other.datatype} image. Both images must have same '
                'datatype.'
            )
        self.ee_image = self.ee_image.addBands(other.ee_image)
        self.bands.extend(other.bands)

        self._check_validity()

    def addBands(self, ee_image, bands):
        "Add new bands to the image"
        self.ee_image = self.ee_image.addBands(ee_image)
        self.bands.extend(bands)

        self._check_validity()

    def cast(self):
        if self.datatype == results.DataType.BYTE:
            self.ee_image = self.ee_image.byte()
        elif self.datatype == results.DataType.UINT16:
            self.ee_image = self.ee_image.uint16()
        elif self.datatype == results.DataType.INT16:
            self.ee_image = self.ee_image.int16()
        elif self.datatype == results.DataType.UINT32:
            self.ee_image = self.ee_image.uint32()
        elif self.datatype == results.DataType.INT32:
            self.ee_image = self.ee_image.int32()
        elif self.datatype == results.DataType.FLOAT32:
            self.ee_image = self.ee_image.float()
        elif self.datatype == results.DataType.FLOAT64:
            self.ee_image = self.ee_image.double()
        else:
            raise GEEImageError(
                f'Unknown datatype {self.datatype}. Datatype '
                'must be supported by GDAL GeoTiff driver.'
            )


def teimage_v1_to_teimage_v2(te_image):
    """Upgrade a version 1 TEImage to TEImageV2"""
    datatype = results.DataType.INT16

    bands = []

    for n, band in enumerate(te_image.band_info):
        # Dump and load each band in order to ensure defaults are added
        band = BandInfoSchema().load(BandInfoSchema().dump(band))
        bands.append(results.Band(**BandInfoSchema().dump(band)))

    image = GEEImage(te_image.image, bands=bands, datatype=datatype)

    return TEImageV2({datatype: image})


# Not using dataclass as not in python 3.6
class TEImageV2():
    "A class to store GEE images and band info for export to cloud storage"

    def __init__(self, images: typing.Dict[str, GEEImage] = {}):
        self.images = images

    def add_image(
        self,
        image: ee.Image,
        bands: typing.List[results.Band],
        datatype: results.DataType = results.DataType.INT16
    ):
        gee_image = GEEImage(image, bands, datatype)

        if datatype not in self.images:
            self.images[datatype] = gee_image
        else:
            self.images[datatype].merge(gee_image)

    def merge(self, other):
        "Merge with another TEImageV2 object"

        for datatype, other_image in other.images.items():
            if datatype in self.images:
                self.images[datatype].merge(other_image)
            else:
                self.images[datatype] = other_image

    def selectBands(self, band_names):
        "Select certain bands from the image(s), dropping all others"

        for image in self.images:
            band_indices = [
                i for i, bi in enumerate(image.bands) if bi.name in band_names
            ]

            if len(band_indices) < 1:
                raise GEEImageError(f'Band(s) "{band_names}" not in image')

            image.bands = [image.bands[i] for i in band_indices]
            image.ee_image = image.ee_image.select(band_indices)

    def setAddToMap(self, band_names=[]):
        "Set the layers that will be added to the map by default"

        for image in self.images:
            for i in range(len(image.bands)):
                if image.bands[i].name in band_names:
                    image.bands[i].add_to_map = True
                else:
                    image.bands[i].add_to_map = False

    def export(
        self,
        geojsons,
        task_name,
        crs,
        logger,
        execution_id=None,
        proj=None,
        filetype=results.RasterFileType.COG
    ):
        "Export layers to cloud storage"

        if not execution_id:
            execution_id = str(random.randint(1000000, 99999999))
        else:
            execution_id = execution_id

        tasks = []

        for datatype, image in self.images.items():
            if not proj:
                proj = image.ee_image.projection()

            image.cast()

            n = 1

            for geojson in geojsons:
                if task_name:
                    out_name = '{}_{}_{}_{}'.format(
                        execution_id, task_name, datatype.value, n
                    )
                else:
                    out_name = '{}_{}_{}'.format(
                        execution_id, datatype.value, n
                    )

                if filetype == results.RasterFileType.COG:
                    as_COG = True
                else:
                    as_COG = False

                export = {
                    'image': image.ee_image,
                    'description': out_name,
                    'fileNamePrefix': out_name,
                    'bucket': BUCKET,
                    'maxPixels': 1e13,
                    'crs': crs,
                    'scale': ee.Number(proj.nominalScale()).getInfo(),
                    'region': get_coords(geojson),
                    'formatOptions': {
                        'cloudOptimized': as_COG
                    }
                }
                t = gee_task(
                    task=ee.batch.Export.image.toCloudStorage(**export),
                    prefix=out_name,
                    logger=logger,
                    metadata={
                        'datatype': datatype,
                        'bands': image.bands
                    }
                )
                tasks.append(t)
                n += 1

        logger.debug("Exporting to cloud storage.")

        output = {}

        for task in tasks:
            task.join()

            if task.metadata['datatype'] in output:
                output[task.metadata['datatype']]['uris'].extend(
                    task.get_uris()
                )
            else:
                output[task.metadata['datatype']] = {
                    'uris': task.get_uris(),
                    'bands': task.metadata['bands']
                }

        rasters = {}

        for datatype, value in output.items():
            uris = [results.URI.Schema().dump(uri) for uri in value['uris']]
            bands = [
                results.Band.Schema().dump(band) for band in value['bands']
            ]

            if len(uris) > 1:
                rasters[datatype.value] = TiledRaster.Schema().load(
                    {
                        'tile_uris': uris,
                        'bands': bands,
                        'datatype': datatype,
                        'filetype': filetype
                    }
                )
            else:
                rasters[datatype.value] = Raster.Schema().load(
                    {
                        'uri': uris[0],
                        'bands': bands,
                        'datatype': datatype,
                        'filetype': filetype
                    }
                )

        gee_results = results.RasterResults(
            name=task_name, rasters=rasters, data={}
        )

        return results.RasterResults.Schema().dump(gee_results)
