import os
import io
import requests
import base64
import time
from multiprocessing import Pool
import io
import uuid
from datetime import datetime

import smart_open
import PIL
import pandas as pd
import numpy as np
import lz4
import boto3
from botocore.exceptions import ClientError
import tqdm
import orjson
import mgzip

# TODO: figure out how to return status codes form lambda functions.
# See the wip_return_error_code branch of the infrastructure repo.
def _raise_for_apigateway_errormessage(response):
    if response is not None and 'errorMessage' in response:
        raise RuntimeError(response['errorMessage'])

def query_dioptra_app(method, path, body=None):
    api_key = os.environ.get('DIOPTRA_API_KEY', None)
    if api_key is None:
        raise RuntimeError('DIOPTRA_API_KEY env var is not set')

    app_endpoint = os.environ.get('DIOPTRA_APP_ENDPOINT', 'https://app.dioptra.ai')

    r = getattr(requests, method.lower())(f'{app_endpoint}{path}', headers={
        'content-type': 'application/json',
        'x-api-key': api_key
    }, json=body)
    r.raise_for_status()

    return r.json()

def download_from_lake(filters, limit=None, order_by=None, desc=None, fields=['*']):
    """
    Downloading metadata from the data lake

    Parameters:
        filters: dioptra style filters to select the data to be queried from
        limit: limit to selected the data
        order_by: field to use to sort the data to control how limit is performed
        desc: whether to order by dec or not
        fields: array of fields to be queried. By default all fields are queried
    """

    return pd.json_normalize(query_dioptra_app(
        'POST',
        '/api/metrics/select',
        {
            'select': ','.join(fields),
            'filters': filters,
            **({'limit': limit} if limit is not None else {}),
            **({'order_by': order_by} if order_by is not None else {}),
            **({'desc': desc} if desc is not None else {})
        }
    ))

def select_datapoints(filters, limit=None, order_by=None, desc=None, fields=['*']):
    """
    Select metadata from the data lake

    Parameters:
        filters: dioptra style filters to select the data to be queried from
        limit: limit to selected the data
        order_by: field to use to sort the data to control how limit is performed
        desc: whether to order by dec or not
        fields: array of fields to be queried. By default all fields are queried
    """

    return pd.DataFrame(query_dioptra_app(
        'POST',
        '/api/datapoints/select',
        {
            'selectColumns': fields,
            'filters': filters,
            **({'limit': limit} if limit is not None else {}),
            **({'orderBy': order_by} if order_by is not None else {}),
            **({'desc': desc} if desc is not None else {})
        }
    ))

def get_predictions(datapoint_ids):
    """
    Get predictions for a set of datapoints

    Parameters:
        datapoint_ids: list of datapoint ids to get predictions for
    """

    return pd.json_normalize(query_dioptra_app(
        'POST',
        '/api/predictions/get',
        {
            'datapointIds': datapoint_ids
        }
    ))

def delete_predictions(prediction_ids):
    """
    Delete predictions from the data lake

    Parameters:
        prediction_ids: list of prediction ids to delete
    """

    return query_dioptra_app(
        'POST',
        '/api/predictions/delete',
        {
            'predictionIds': prediction_ids
        }
    )

def get_groundtruths(datapoint_ids):
    """
    Get groundtruths for a set of datapoints

    Parameters:
        datapoint_ids: list of datapoint ids to get groundtruths for
    """

    return pd.json_normalize(query_dioptra_app(
        'POST',
        '/api/groundtruths/get',
        {
            'datapointIds': datapoint_ids
        }
    ))

def delete_groundtruths(groundtruth_ids):
    """
    Delete groundtruths from the data lake

    Parameters:
        groundtruth_ids: list of groundtruth ids to delete
    """

    return query_dioptra_app(
        'POST',
        '/api/groundtruths/delete',
        {
            'groundtruthIds': groundtruth_ids
        }
    )

def delete_datapoints(filters, limit=None, order_by=None, desc=None):
    """
    Delete metadata from the data lake

    Parameters:
        filters: dioptra style filters to select the data to be deleted
        limit: limit to selected the data
        order_by: field to use to sort the data to control how limit is performed
        desc: whether to order by dec or not
    """

    return query_dioptra_app(
        'POST',
        '/api/datapoints/delete',
        {
            'filters': filters
        }
    )

def delete_from_lake(filters, limit=None, order_by=None, desc=None):
    """
    Delete metadata from the data lake

    Parameters:
        filters: dioptra style filters to select the data to be deleted
        limit: limit to selected the data
        order_by: field to use to sort the data to control how limit is performed
        desc: whether to order by dec or not
    """

    return query_dioptra_app(
        'POST',
        '/api/metrics/delete',
        {
            'filters': filters,
            **({'limit': limit} if limit is not None else {}),
            **({'order_by': order_by} if order_by is not None else {}),
            **({'desc': desc} if desc is not None else {})
        }
    )

def upload_to_lake(records):
    """
    Uploading metadata to the data lake

    Parameters:
        records: array of dipotra style records. See teh doc for accepted formats
    """

    api_key = os.environ.get('DIOPTRA_API_KEY', None)
    if api_key is None:
        raise RuntimeError('DIOPTRA_API_KEY env var is not set')

    api_endpoint = os.environ.get('DIOPTRA_API_ENDPOINT', 'https://api.dioptra.ai/events')

    try:
        r = requests.post(api_endpoint, headers={
            'content-type': 'application/json',
            'x-api-key': api_key
        }, json={
            'records': records
        })
        r.raise_for_status()
        response = r.json()
        _raise_for_apigateway_errormessage(response)
    except requests.exceptions.RequestException as err:
        print('There was an error uploading to the lake ...')
        raise err

    return response

def upload_to_lake_from_s3(bucket_name, object_name, storage_type='s3'):
    """
    Uploading metadata to the data lake from an s3 bucket
    The data should be a new line delimited JSON and can be compressed with Gzip
    Will generate a signed URL and pass it to the dioptra ingestion
    boto3 credentials should be configured

    Parameters:
        bucket_name: name of the bucket where the records are
        object_name: name of the key in the bucket where the records are
    """

    api_key = os.environ.get('DIOPTRA_API_KEY', None)
    if api_key is None:
        raise RuntimeError('DIOPTRA_API_KEY env var is not set')

    api_endpoint = os.environ.get('DIOPTRA_API_ENDPOINT', 'https://api.dioptra.ai/events')

    if storage_type == 's3':
        signed_url = _generate_s3_signed_url(bucket_name, object_name)

    try:
        r = requests.post(api_endpoint, headers={
            'content-type': 'application/json',
            'x-api-key': api_key
        }, json={
            'urls': [signed_url]
        })
        r.raise_for_status()
        response = r.json()
        _raise_for_apigateway_errormessage(response)
    except requests.exceptions.RequestException as err:
        print('There was an error uploading to the lake ...')
        raise err

    return response

def wait_for_upload(upload_id):
    """
    Wait for the upload status to be SUCCEEDED | FAILED | TIMED_OUT | ABORTED
    """
    if upload_id is None:
        raise RuntimeError('upload_id must not be None')

    api_key = os.environ.get('DIOPTRA_API_KEY', None)
    if api_key is None:
        raise RuntimeError('DIOPTRA_API_KEY env var is not set')

    app_endpoint = os.environ.get('DIOPTRA_APP_ENDPOINT', 'https://app.dioptra.ai')
    sleepTimeSecs = 1
    totalSleepTimeSecs = 0
    try:
        while True:
            if totalSleepTimeSecs > 900:
                raise RuntimeError('Timed out waiting for the upload to finish.')

            r = requests.get(f'{app_endpoint}/api/ingestion/executions/{upload_id}', headers={
                'content-type': 'application/json',
                'x-api-key': api_key
            })
            r.raise_for_status()
            upload = r.json()
            if upload['status'] in ['SUCCEEDED', 'FAILED', 'TIMED_OUT', 'ABORTED']:
                return upload
            else:
                time.sleep(sleepTimeSecs)
                totalSleepTimeSecs += sleepTimeSecs
                sleepTimeSecs = min(sleepTimeSecs * 2, 60)
    except requests.exceptions.RequestException as err:
        print('There was an error waiting for the upload to finish ...')
        raise err

def _generate_s3_signed_url(bucket_name, object_name):

    s3_client = boto3.client('s3')
    try:
        response = s3_client.generate_presigned_url(
            'get_object', Params={'Bucket': bucket_name, 'Key': object_name}, ExpiresIn=3600)
    except ClientError as err:
        print('There was an error getting a signed URL ...')
        raise err

    return response

def _list_dataset_metadata():
    """
    List all the dataset metadata

    """

    return query_dioptra_app('GET', '/api/dataset')

def _list_miner_metadata():
    """
    List all the miner metadata

    """

    return query_dioptra_app('GET', '/api/miner')


def _encode_np_array(np_array):
    """
    Encode and compress a np array

    Parameters:
        np_array: the np array to be encoded

    """
    if not isinstance(np_array, np.ndarray):
        raise RuntimeError('Can only encode numpy arrays')

    bytes_buffer = io.BytesIO()
    np.save(bytes_buffer, np_array)

    return base64.b64encode(
        lz4.frame.compress(
            bytes_buffer.getvalue(),
            compression_level=lz4.frame.COMPRESSIONLEVEL_MAX
        )).decode('ascii')

def _decode_to_np_array(value):
    """
    Decode a compress a np array

    Parameters:
        value: the string containing the np array

    """
    decoded_bytes = lz4.frame.decompress(base64.b64decode(value))
    return np.load(io.BytesIO(decoded_bytes), allow_pickle=True)


def _format_groundtruth(groundtruth, gt_type, class_names=None):
    """
    Utility formatting the groundtruth field according to the model type

    Parameters:
        groundtruth: the groundtruth field
        gt_type: the type of gt. Supported types CLASSIFIER, SEMANTIC_SEGMENTATION
        class_names: a list of class names. If the groundtruth contains indexes, it will be used to convert them to names

    """
    if gt_type == 'CLASSIFIER':
        if isinstance(groundtruth, int) and class_names is not None:
            class_name = class_names[groundtruth]
        else:
            class_name = groundtruth
        return {
            'class_name': class_name
        }
    if gt_type == 'SEMANTIC_SEGMENTATION':
        if isinstance(groundtruth, PIL.Image.Image):
            gt_array = np.array(groundtruth)
        if isinstance(groundtruth, np.ndarray):
            gt_array = groundtruth
        if isinstance(groundtruth, list):
            gt_array = np.array(groundtruth)
        return {
            'encoded_segmentation_class_mask': _encode_np_array(gt_array),
            **({'class_names': class_names} if class_names is not None else {})
        }


def upload_image_dataset(
    dataset, dataset_type,
    image_field=None,
    groundtruth_field=None,
    datapoints_metadata=None,
    class_names=None, dataset_tags=None,
    max_batch_size=1000, num_workers=1):
    """
    Upload an image dataset to Dioptra.
    The images will be uploaded to a bucket specified with DIOPTRA_UPLOAD_BUCKET.
    The metadata will be uploaded to Dioptra and point to this bucket

    Parameters:
        dataset: the dataset to upload. Should be iteratable.
        image_field: the name of the field containing the image. Should be in PIL format
        dataset_type: the type of data, Supported CLASSIFIER, SEMANTIC_SEGMENTATION
        groundtruth_field: the name of the field containing the groundtruth
        datapoints_metadata:
            a list of metadata to be added to each datapoint.
            should already be formatted.
            if it contains a 'image_id' field, this will be used to name the images in s3
        class_names: a list of class names to convert indexes in the groundtruth to class names
        class_names: a dict of tags to be added to the entire dataset
        max_batch_size: the maximum batch uplodd size
        num_workers: number of parallel workers to upload the images to the bucket

    """


    def _img_upload(payload):
        image_bytes, img_url = payload[0], payload[1]
        with smart_open.open(img_url, 'wb') as file:
            file.write(image_bytes)

    def _upload_data(storage_type, bucket, prefix, dataset_metadata, img_payload, num_workers):
        with Pool(num_workers) as my_pool:
            my_pool.map(_img_upload, img_payload)

        file_name = os.path.join(
            prefix, 'logs', f'{str(uuid.uuid4())}_{datetime.utcnow().isoformat()}.ndjson')
        logs_url = f'{storage_type}://{os.path.join(bucket, file_name)}'

        payload = b'\n'.join(
            [orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)for record in dataset_metadata])
        compressed_payload = mgzip.compress(payload, compresslevel=2)

        with smart_open.open(logs_url, 'wb') as file:
            file.write(compressed_payload)

        return upload_to_lake_from_s3(bucket, file_name, storage_type)

    def _build_img_url(storage_type, bucket, prefix, image_id, pil_image):
        img_format = pil_image.format
        root, ext = os.path.splitext(image_id)
        if not ext:
            image_id += f'.{img_format.lower()}'
        return f'{storage_type}://{os.path.join(bucket, prefix, "images", image_id)}'

    api_key = os.environ.get('DIOPTRA_API_KEY', None)
    if api_key is None:
        raise RuntimeError('DIOPTRA_API_KEY env var is not set')

    s3_bucket = os.environ.get('DIOPTRA_UPLOAD_BUCKET', None)

    if s3_bucket is None:
        raise RuntimeError('DIOPTRA_UPLOAD_BUCKET env var is not set')

    s3_prefix_bucket = os.environ.get('DIOPTRA_UPLOAD_PREFIX', '')
    storage_type = os.environ.get('DIOPTRA_UPLOAD_BUCKET_TYPE', 's3')

    dataset_metadata = []
    img_payload = []
    upload_ids = []

    for index, row in tqdm.tqdm(enumerate(dataset), desc='uploading dataset...'):

        if image_field is not None and image_field not in row:
            raise RuntimeError(f'{image_field} is not in the dataset')
        if groundtruth_field is not None and groundtruth_field not in row:
            raise RuntimeError(f'{groundtruth_field} is not in the dataset')

        if image_field is not None:
            pil_img = row[image_field]
            if not isinstance(pil_img, PIL.Image.Image):
                raise Exception('the image provided was not a PIL image')
            img_width, img_height = pil_img.size
            image_id = str(index)

        my_metadata = None
        if datapoints_metadata is not None:
            my_metadata = datapoints_metadata[index]
            if 'image_id' in my_metadata:
                image_id = my_metadata['image_id']

        img_url = _build_img_url(storage_type, s3_bucket, s3_prefix_bucket, image_id, pil_img)

        in_mem_file = io.BytesIO()
        pil_img.save(in_mem_file, format=pil_img.format)

        datapoint_tags = {}
        image_metadata = {}

        if  my_metadata is not None and 'tags' in my_metadata:
            datapoint_tags.update(my_metadata['tags'])
        if dataset_tags is not None:
            datapoint_tags.update(dataset_tags)
        if  my_metadata is not None and 'image_metadata' in my_metadata:
            image_metadata.update(my_metadata['image_metadata'])

        image_metadata.update({
            'uri': img_url,
            'width': img_width,
            'height': img_height
        })

        img_payload.append((in_mem_file.getvalue(), img_url))
        dataset_metadata.append({
            'image_metadata': image_metadata,
            'model_type': dataset_type,
            **({
                'groundtruth': _format_groundtruth(row[groundtruth_field], dataset_type, class_names)
            } if groundtruth_field is not None else {}),
            **(my_metadata if my_metadata is not None else {}),
            **({'tags': datapoint_tags} if len(datapoint_tags) > 0 else {})
        })

        if len(img_payload) > max_batch_size:
            upload_ids.append(_upload_data(
                storage_type, s3_bucket, s3_prefix_bucket, dataset_metadata, img_payload, num_workers))
            dataset_metadata = []
            img_payload = []

    if len(img_payload) > 0:
        upload_ids.append(_upload_data(
            storage_type, s3_bucket, s3_prefix_bucket, dataset_metadata, img_payload, num_workers))

    return upload_ids
