import requests
import json
import boto3
import os
import pathlib
import random
import time
from datetime import datetime
from uuid import uuid4
from aws_requests_auth.aws_auth import AWSRequestsAuth
#from BatchLogger import BatchLogger
from pathlib import Path
#from Utils import Utils


class GMAI:
    def __init__(self, task, access_key_id, secret_access_key):
        self.task = task
        self._bucket = 'web-rendering'
        self._region = 'eu-west-1'
        self._api_service = 'execute-api'

        self._api_host = 'fp13kcc0ug.execute-api.eu-west-1.amazonaws.com'
        self._invoke_url = 'https://' + self._api_host
        self._folder = 'alfalfa'
        self._alfalfa_rendering_endpoint = '/dev/start_rendering'
        self._parameters_extraction_endpoint = '/dev/parameters_extraction'

        self._auth = AWSRequestsAuth(aws_access_key=access_key_id,
                                     aws_secret_access_key=secret_access_key,
                                     aws_host=self._api_host,
                                     aws_region=self._region,
                                     aws_service=self._api_service)
        self._s3_client = boto3.client('s3', region_name=self._region, aws_access_key_id=access_key_id,
                                       aws_secret_access_key=secret_access_key)
        self._batch_client = boto3.client(service_name='batch', region_name=self._region,
                                          endpoint_url='https://batch.' + self._region + '.amazonaws.com')
        self._cloudwatch_client = boto3.client(service_name='logs', region_name=self._region,
                                               endpoint_url='https://logs.' + self._region + '.amazonaws.com')
        self._sf_client = boto3.client('stepfunctions', region_name=self._region, aws_access_key_id=access_key_id,
                                       aws_secret_access_key=secret_access_key)
        #self._utils = Utils(self._s3_client, self._bucket, self._folder)

    '''def query(self, image_path):
        folder_path, file_name = os.path.split(image_path)
        s3_path = self._folder + '/data/' + file_name
        self._s3_client.upload_file(image_path, self._bucket, s3_path)

        test_dict = {"image_path": s3_path, "file_name": file_name, "bucket": self._bucket}

        response = requests.post(self._url, auth=self._auth, json=json.dumps(test_dict))

        if folder_path == "":
            output_file_name = "output.jpg"
        else:
            output_file_name = "\\output.jpg"
        output_file_path = folder_path + output_file_name
        self._s3_client.download_file(self._bucket, json.loads(response.text)["s3_path"], output_file_path)

        return output_file_path'''

    def get_params(self):
        if self.task == 'alfalfa':
            print(f"Possible parameters to change for {self.task}:")
            print('min_distance, default value: 0.15 m')
            print('min_density, default value: 0.4')
            print('max_density, default value: 1.0')
            print('min_num_branches, default value: 3')
            print('max_num_branches, default value: 20')
        else:
            print('Wrong initialization,task name not supported. Possible values: alfalfa')

    def upload_blend_file(self, blend_file_path, folder_path=None):
        if folder_path is None and pathlib.PurePath(blend_file_path).suffix.lower() != '.blend':
            raise Exception("Wrong extention, this is not a blender file")

        uuid = str(uuid4())
        folder = uuid + '/blender_file/'
        bucket, folder_on_s3, blender_file_name = self.upload_file_to_s3(blend_file_path, folder)
        if folder_path is not None:
            self.upload_folder_to_s3(folder_path, folder)

        params = {
            'bucket': bucket,
            'folder': folder_on_s3,
            'blender_file_name': blender_file_name
        }
        '''params = {
            'bucket': 'web-rendering',
            'folder': 'blender_files/alfalfa/',
            'blender_file_name': 'alfalfa.blend'
        }'''
        print("All files uploaded, extracting parameters")
        response = requests.post(self._invoke_url + self._parameters_extraction_endpoint, auth=self._auth, json=json.dumps(params))

        if response.ok:
            response_body = response.json()
            print('Available parameters:')
            print(json.dumps(response_body, indent=2))
        else:
            print(response.text)
            print(response)
            print('Error while sending request')
            print(response.json())

    def render_alfalfa(self, nr_imgs_to_render, nr_return_imgs, min_distance=0.15, min_density=0.4, max_density=1.0,
                       min_num_branches=3, max_num_branches=20):
        params = {
            'parameters': {
                'start_scene': 0,
                'end_scene': nr_imgs_to_render - 1,
                #'task': self.task,
                'min_distance': min_distance,
                'min_density': min_density,
                'max_density': max_density,
                'min_num_branches': min_num_branches,
                'max_num_branches': max_num_branches
            }
        }

        response = requests.post(self._invoke_url + self._alfalfa_rendering_endpoint, auth=self._auth, json=json.dumps(params))

        if response.ok:
            response_body = response.json()
            job_s3_bucket = response_body['s3Bucket']
            s3_prefix = response_body['s3Prefix']
            job_name = response_body['jobName']
            job_uuid = response_body['uuid']
            execution_arn = response_body['executionArn']

            #batch_logger = BatchLogger(self._cloudwatch_client, self._batch_client, self._sf_client)

            status = self.print_step_func_status(execution_arn)
            if status == 'SUCCEEDED':
                self._download_results(s3_prefix=s3_prefix, nr_return_imgs=nr_return_imgs)

        else:
            print(response.text)
            print(response)
            print('Error while sending request')
            print(response.json())

    '''def render(self, config):
        params = {
            'parameters': config
        }
        response = requests.post(self._url, auth=self._auth, json=json.dumps(params))

        if response.ok:
            response_body = response.json()
            # print(response_body)
            job_id = response_body['jobId']
            job_name = response_body['jobName']
            job_uuid = response_body['uuid']

            batch_logger = BatchLogger(self._cloudwatch_client, self._batch_client)

            status = batch_logger.print_status(job_name, job_id)
            if status == 'SUCCEEDED':
                # TODO change folder name to job_id
                self._download_results(folder=job_uuid)

        else:
            print('Error while sending request')
            print(response.json())

        # print(response)
        # print(response.text)'''

    def _list_files_in_s3(self, s3_prefix=''):
        params = {
            "Bucket": self._bucket,
            "Prefix": s3_prefix
        }

        result = self._s3_client.list_objects_v2(**params)

        if 'Contents' not in result:
            print(result)
            return []

        file_list = []
        for key in result['Contents']:
            file_list.append(key['Key'])

        while result['IsTruncated']:
            continuation_key = result['NextContinuationToken']
            result = self._s3_client.list_objects_v2(ContinuationToken=continuation_key, **params)
            for key in result['Contents']:
                file_list.append(key['Key'])
        return file_list

    def _download_results(self, nr_return_imgs, folder=None, s3_prefix=None):
        files_names_list = self._list_files_in_s3(s3_prefix)

        tsne_file = next(filter(lambda file_name: 'tsne.jpg' in file_name, files_names_list))

        files_names_list.remove(tsne_file)
        files_names_list = random.sample(files_names_list, nr_return_imgs)
        files_names_list.append(tsne_file)

        files_count = len(files_names_list)

        for idx, file in enumerate(files_names_list):
            print(f'Downloading file {idx}/{files_count-1}')
            # skip folders
            if file[-1] == "/":
                continue

            # create nested directory structure
            obj_path = os.path.dirname(file)
            Path(obj_path).mkdir(parents=True, exist_ok=True)

            # save file with full path locally
            self._s3_client.download_file(self._bucket, file, file)

        print(f"Files downloaded to folder: {obj_path}")

    def print_step_func_status(self, execution_arn):
        spin = ['.', '..', '...']
        spinner = 0

        while True:
            time.sleep(5)
            sf_response = self._sf_client.describe_execution(executionArn=execution_arn)
            status = sf_response['status']
            print(status)
            if status == 'SUCCEEDED':
                print(f"{'=' * 80}")
                print(f'End of the execution: {status}')
                break
            elif status in ['FAILED', 'TIMED_OUT', 'ABORTED']:
                print(f"{'=' * 80}")
                print(f'End of the execution: {status}')
                break
            elif status == 'RUNNING':
                print(f'Execution still RUNNING {spin[spinner]}')
                spinner = (spinner + 1) % len(spin)
            else:
                print(f'Unexpected status of the execution: {status}')
                break

        return status

    def print_status(self, job_name, job_id):
        spin = ['.', '..', '...']
        spinner = 0
        start_time = 0
        running = False

        while True:
            time.sleep(5)
            describe_jobs_response = self._batch_client.describe_jobs(jobs=[job_id])
            status = describe_jobs_response['jobs'][0]['status']
            print(status)
            if status == 'SUCCEEDED':
                print(f"{'=' * 80}")
                print(f'Job [{job_id}] {status}')
                break
            elif status == 'FAILED':
                print(f"{'=' * 80}")
                print(f'Job [{job_id}] {status}')
                break
            elif status == 'RUNNING':
                log_stream_name = self._get_log_stream(self._log_group_name, job_name, job_id)
                if not running and log_stream_name:
                    running = True
                    print(f'\rJob [{job_name} - {job_id}] is RUNNING.')
                    print(f"Output [{log_stream_name}]:\n {'=' * 80}")
                if log_stream_name:
                    start_time = self.print_logs(self._log_group_name, log_stream_name, start_time) + 1
            else:
                print(f'\rJob [{job_name} - {job_id}] is {status}  {spin[spinner]}')
                spinner = (spinner + 1) % len(spin)

        return status

    def print_logs(self, log_group_name, log_stream_name, start_time):
        kwargs = {'logGroupName': log_group_name,
                  'logStreamName': log_stream_name,
                  'startTime': start_time,
                  'startFromHead': True}

        last_timestamp = 0
        while True:
            log_events = self._cloudwatch_client.get_log_events(**kwargs)

            for event in log_events['events']:
                last_timestamp = event['timestamp']
                timestamp = datetime.utcfromtimestamp(last_timestamp / 1000.0).isoformat()
                print('[%s] %s' % ((timestamp + ".000")[:23] + 'Z', event['message']))

            next_token = log_events['nextForwardToken']
            if next_token and kwargs.get('nextToken') != next_token:
                kwargs['nextToken'] = next_token
            else:
                break
        return last_timestamp

    def _get_log_stream(self, log_group_name, job_name, job_id):
        response = self._cloudwatch_client.describe_log_streams(
            logGroupName=log_group_name,
            logStreamNamePrefix=job_name + '/' + job_id
        )
        log_streams = response['logStreams']
        if not log_streams:
            return ''
        else:
            return log_streams[0]['logStreamName']

    def upload_file_to_s3(self, local_file_path, folder, file_path_on_s3=None):
        if file_path_on_s3 is None:
            file_name = pathlib.PurePath(local_file_path).name
        else:
            file_name = file_path_on_s3

        # s3_folder_path = self._s3_folder + '/' + folder
        s3_folder_path = folder
        s3_path = s3_folder_path + file_name
        print(f"Uploading {file_name}")
        self._s3_client.upload_file(local_file_path, self._bucket, s3_path)

        return self._bucket, s3_folder_path, file_name

    def upload_folder_to_s3(self, local_folder_path, folder):
        path = pathlib.Path(local_folder_path)
        files_to_upload = [item for item in path.rglob("*") if not item.is_dir()]
        for file in files_to_upload:
            relative_file_path = str(file.relative_to(path.parent).as_posix())
            self.upload_file_to_s3(str(file), folder, relative_file_path)

