import logging
import datetime
import json
import os
import gzip
import numpy as np
import threading
import random

from tqdm import tqdm
import boto3
from botocore import UNSIGNED
from botocore.client import Config
import uuid

import dioptra
from dioptra.client import Client
from dioptra.supported_types import SupportedTypes
from dioptra.utils import (
    validate_tags,
    validate_embeddings,
    validate_features,
    validate_annotations,
    validate_confidence,
    validate_image_metadata,
    validate_timestamp,
    validate_text_metadata,
    validate_text,
    validate_audio_metadata,
    validate_model_type,
    validate_input_data
)

s3_client = boto3.client('s3', config=Config(signature_version=UNSIGNED))

class FileProgress(object):
    def __init__(self, filenames):
        self._size = 0

        for filename in filenames:
            self._size += float(os.path.getsize(filename))

        self._lock = threading.Lock()
        self._pbar = tqdm(
            total=self._size,
            unit='B',
            desc='Uploading to Dioptra'
        )

    def __call__(self, bytes_amount):
        with self._lock:
            self._pbar.update(bytes_amount)

BATCH_INGEST_MAX_SIZE = 10000

class Logger:
    def __init__(
        self,
        api_key,
        endpoint_url='https://api.dioptra.ai',
        synchronous_mode=False,
        max_workers=5,
        batch_size=20,
        queue_size=100
    ):
        self.api_key = api_key
        self.endpoint_url = endpoint_url
        self.event_url = endpoint_url + '/events'
        self._headers = {
            'x-api-key': api_key,
            'Content-Type': 'application/json'
        }
        # TODO: reenable setting synchronous_mode
        self.client = Client(max_workers, queue_size, synchronous_mode=True)
        self.batch = []
        self.batch_size = batch_size

    def log(
        self,
        request_id,
        timestamp=None,
        dataset_id=None,
        benchmark_id=None,
        model_id=None,
        model_version=None,
        model_type=None,
        groundtruth=None,
        prediction=None,
        confidence=None,
        features=None,
        embeddings=None,
        input_data=None,
        image_metadata=None,
        text=None,
        text_metadata=None,
        tags=None,
        audio_metadata=None
    ):
        """
        The log interface sends a set of events to the dioptra api endpoint.
        Yet, the events are not commited yet.
        As long as a event is not commited, it won't be available for analysis.
        Once a committed event with the same `request_id` is sent,
        it will be aggregated with previous events in a session that share the same `request_id`
        A session is 1 day long.

        """

        payload = self.package_payload(
            model_id,
            dataset_id,
            benchmark_id,
            model_version,
            model_type,
            timestamp,
            request_id,
            groundtruth,
            prediction,
            confidence,
            features,
            embeddings,
            input_data,
            image_metadata,
            text,
            text_metadata,
            tags,
            audio_metadata,
            False
        )

        self.client.call(
            payload=payload,
            endpoint_url=self.event_url,
            headers=self._headers)

    def commit(
        self,
        request_id,
        timestamp=None,
        model_id=None,
        dataset_id=None,
        benchmark_id=None,
        model_version=None,
        model_type=None,
        groundtruth=None,
        prediction=None,
        confidence=None,
        features=None,
        embeddings=None,
        input_data=None,
        image_metadata=None,
        text=None,
        text_metadata=None,
        tags=None,
        audio_metadata=None
    ):
        """
        The commit interface sends a set of events to the dioptra api endpoint and commits them.
        A committed event triggers the aggregation of all previous events with the same `request_id`
        in the same session. A seesion is 1 day long.
        Once a committed all events are available for analysis in the dioptra UI.

        """

        payload = self.package_payload(
            model_id,
            dataset_id,
            benchmark_id,
            model_version,
            model_type,
            timestamp,
            request_id,
            groundtruth,
            prediction,
            confidence,
            features,
            embeddings,
            input_data,
            image_metadata,
            text,
            text_metadata,
            tags,
            audio_metadata,
            True
        )

        self.client.call(
            payload=payload,
            endpoint_url=self.event_url,
            headers=self._headers)

    def add_to_batch_log(
        self,
        request_id,
        timestamp=None,
        dataset_id=None,
        benchmark_id=None,
        model_id=None,
        model_version=None,
        model_type=None,
        groundtruth=None,
        prediction=None,
        confidence=None,
        features=None,
        embeddings=None,
        input_data=None,
        image_metadata=None,
        text=None,
        text_metadata=None,
        tags=None,
        audio_metadata=None
    ):

        payload = self.package_payload(
            model_id,
            dataset_id,
            benchmark_id,
            model_version,
            model_type,
            timestamp,
            request_id,
            groundtruth,
            prediction,
            confidence,
            features,
            embeddings,
            input_data,
            image_metadata,
            text,
            text_metadata,
            tags,
            audio_metadata,
            False
        )

        self.batch.extend(payload)


    def add_to_batch_commit(
        self,
        request_id,
        timestamp=None,
        dataset_id=None,
        benchmark_id=None,
        model_id=None,
        model_version=None,
        model_type=None,
        groundtruth=None,
        prediction=None,
        confidence=None,
        features=None,
        embeddings=None,
        input_data=None,
        image_metadata=None,
        text=None,
        text_metadata=None,
        tags=None,
        audio_metadata=None
    ):

        payload = self.package_payload(
            model_id,
            dataset_id,
            benchmark_id,
            model_version,
            model_type,
            timestamp,
            request_id,
            groundtruth,
            prediction,
            confidence,
            features,
            embeddings,
            input_data,
            image_metadata,
            text,
            text_metadata,
            tags,
            audio_metadata,
            True
        )

        self.batch.extend(payload)


    def commit_ndjson_file(self, ndjson_filename, validate_sample=0.01, batch_bucket='dioptra-batch-input-prod'):
        """
        Loading the config from file

        """
        events = []
        with open(ndjson_filename, 'r') as ndjson_file:
            lines = ndjson_file.read().split('\n')
            all_lines = np.tile(lines, 1).tolist() # multiples of 100 events
            for line in tqdm(
                all_lines, 
                desc=f'Validating datapoints (sampling: {100 * validate_sample}%)',
                unit='datapoints'
            ):
                parsed_line = json.loads(line)
                events.extend(self.package_payload(**parsed_line, validate_sample=validate_sample))

        # Split the ingestion into files of roughly BATCH_INGEST_MAX_SIZE events.
        # This is allow parallelizing with lambda functions.
        event_batches = np.array_split(events, max(BATCH_INGEST_MAX_SIZE, len(events)) / BATCH_INGEST_MAX_SIZE)
        batch_keys = [str(uuid.uuid4()) + '.ndjson.gz' for _ in event_batches]
        batch_tmp_filepaths = [f'/tmp/dioptra/{key}' for key in batch_keys]

        for i, event_batch in enumerate(tqdm(
            event_batches, 
            desc=f'Compressing dataset chunks',
            unit='chunk'
        )):
            events_payload = '\n'.join([json.dumps(e) for e in event_batch]).encode()

            tmp_filepath = batch_tmp_filepaths[i]
            os.makedirs(os.path.dirname(tmp_filepath), exist_ok=True)
            tmp_file = open(tmp_filepath, 'wb')
            tmp_file.write(gzip.compress(events_payload))

        file_progress = FileProgress(batch_tmp_filepaths)
        transfer = boto3.s3.transfer.S3Transfer(s3_client)

        for i, tmp_filepath in enumerate(batch_tmp_filepaths):
            try:
                key = batch_keys[i]
                transfer.upload_file(
                    tmp_filepath, batch_bucket, key,
                    callback=file_progress,
                    extra_args={'Metadata': {'x-api-key': self.api_key}}
                )
            except Exception as e:
                print(repr(e))
            finally:
                os.remove(tmp_filepath)

    def package_payload(
        self,
        model_id=None,
        dataset_id=None,
        benchmark_id=None,
        model_version=None,
        model_type=None,
        timestamp=None,
        request_id=None,
        groundtruth=None,
        prediction=None,
        confidence=None,
        features=None,
        embeddings=None,
        input_data=None,
        image_metadata=None,
        text=None,
        text_metadata=None,
        tags=None,
        audio_metadata=None,
        committed=None,
        validate_sample=1
    ):
        skip_validation = random.random() > validate_sample

        payload = {
            'committed': committed,
            'api_version': dioptra.__version__
        }

        if request_id:
            payload['request_id'] = request_id

        if model_id:
            payload['model_id'] = model_id

        if dataset_id:
            payload['dataset_id'] = dataset_id

        if benchmark_id:
            payload['benchmark_id'] = benchmark_id

        if model_version:
            payload['model_version'] = model_version

        if model_type:
            if isinstance(model_type, str):
                model_type = SupportedTypes[model_type]

            if validate_model_type(model_type):
                payload['model_type'] = model_type.model_type.value
                payload['input_type'] = model_type.input_type.value
            else:
                logging.warning('Invalid model_type. Skipping datapoint...')
                return []

        if timestamp and validate_timestamp(timestamp):
            payload['timestamp'] = timestamp.isoformat()
        else:
            logging.info('Invalid timestamp. Replacing with current time...')
            payload['timestamp'] = datetime.datetime.utcnow().isoformat()

        if groundtruth:
            if model_type:
                if skip_validation or validate_model_type(model_type):
                    if skip_validation or validate_annotations(groundtruth, model_type):
                        payload['groundtruth'] = groundtruth
                    else:
                        logging.warning('Invalid groundtruth. Skipping datapoint...')
                        return []
                else:
                    logging.warning('Invalid model_type. Skipping datapoint...')
                    return []
            else:
                logging.warning('No model_type. Skipping datapoint...')
                return []

        if prediction:
            if model_type:
                if skip_validation or validate_model_type(model_type):
                    if skip_validation or validate_annotations(prediction, model_type):
                        payload['prediction'] = prediction
                    else:
                        logging.warning('Invalid prediction. Skipping datapoint...')
                        return []
                else:
                    logging.warning('Invalid model_type. Skipping datapoint...')
                    return []
            else:
                logging.warning('No model_type defined. Skipping datapoint...')
                return []

        if confidence:
            if skip_validation or validate_confidence(confidence):
                payload['confidence'] = confidence
            else:
                logging.warning('Invalid confidence. Skipping datapoint...')
                return []

        if features:
            if skip_validation or validate_features(features):
                payload['features'] = features
            else:
                logging.warning('Invalid features. Skipping datapoint...')
                return []

        if image_metadata:
            if skip_validation or validate_image_metadata(image_metadata):
                payload['image_metadata'] = image_metadata
            else:
                logging.warning('Invalid image_metadata. Skipping datapoint...')
                return []

        if text:
            if skip_validation or validate_text(text):
                payload['text'] = text
            else:
                logging.warning('Invalid text. Skipping datapoint...')
                return []

        if text_metadata:
            if skip_validation or validate_text_metadata(text_metadata):
                payload['text_metadata'] = text_metadata
            else:
                logging.warning('Invalid text_metadata. Skipping datapoint...')
                return []

        if embeddings:
            if skip_validation or validate_embeddings(embeddings):
                payload['embeddings'] = embeddings
            else:
                logging.warning('Invalid embeddings. Skipping datapoint...')
                return []

        if tags:
            if skip_validation or validate_tags(tags):
                payload['tags'] = tags
            else:
                logging.warning('Invalid tags. Skipping datapoint...')
                return []

        if audio_metadata:
            if skip_validation or validate_audio_metadata(audio_metadata):
                payload['audio_metadata'] = audio_metadata
            else:
                logging.warning('Invalid audio_metadata. Skipping datapoint...')
                return []

        if input_data:
            if model_type:
                if skip_validation or validate_model_type(model_type):
                    if skip_validation or validate_input_data(input_data, model_type):
                        payload['input_data'] = input_data
                    else:
                        logging.warning('Invalid input_data. Skipping datapoint...')
                        return []
                else:
                    logging.warning('Invalid model_type. Skipping datapoint...')
                    return []
            else:
                logging.warning('No model model_type defined. Skipping datapoint...')
                return []

        return [payload]
