#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


from timeit import default_timer as timer
import logging
import boto3
from botocore.exceptions import ClientError

import time
import urllib.request
import zipfile
import io


MAX_LOG_URLS_RETRIEVAL_RETRIES = 5


LOG_DIR_PREFIX = "var/log/eb-docker/containers/eb-current-app"


def process_zip_file(file_obj, worker_instance_id):

    relevant_lines_for_log = []
    with zipfile.ZipFile(file_obj) as zip_ref:

        log_files = filter(lambda info: info.filename.startswith(LOG_DIR_PREFIX), zip_ref.infolist())
        sorted_log_files_list = sorted(log_files, key=lambda x: x.date_time, reverse=True)

        # Only examine one log per zip file
        for info in sorted_log_files_list[:1]:
            with zip_ref.open(info) as log_fh:
                log_lines = log_fh.readlines()

                with open(worker_instance_id + ".log", "wb") as output_fh:
                    for line in log_lines:
                        output_fh.write(line)

                for line in log_lines:
                    line_string = line.decode('UTF-8').strip()
                    if line_string.startswith("Posted to: /worker/scan-sha1"):

                        relevant_lines_for_log.append(line_string)

    return relevant_lines_for_log


def get_eb_worker_logs(eb_environment_id):

    eb_client = boto3.client('elasticbeanstalk', region_name='us-east-2')

    try:
        msg = eb_client.request_environment_info(
                EnvironmentId=eb_environment_id,
                InfoType='bundle',
        )

        print("First message:", msg)

        for i in range(MAX_LOG_URLS_RETRIEVAL_RETRIES):
            msg2 = eb_client.retrieve_environment_info(
                    EnvironmentId=eb_environment_id,
                    InfoType='bundle',
            )

            environment_info_list = msg2.get("EnvironmentInfo", [])

            if environment_info_list:
                log_timestamp_url_tuples_by_instance_id = {}
                for log_item in environment_info_list:
                    s3_url = log_item['Message']

                    log_timestamp = log_item['SampleTimestamp']
                    ec2_instance_id = log_item['Ec2InstanceId']

                    log_timestamp_url_tuples_by_instance_id.setdefault(ec2_instance_id, []).append((log_timestamp, s3_url, ec2_instance_id))

                log_timestamp_url_tuples = list(map(lambda x: x[0], sorted(log_timestamp_url_tuples_by_instance_id.values(), key=lambda x: x[0], reverse=True)))
                print("Log URL count:", len(log_timestamp_url_tuples))

                return log_timestamp_url_tuples

            else:
                print("Environment info was empty on iteration %d. Sleeping..." % i)
                time.sleep(5)

    except ClientError as e:
        logging.error(e)
        return None


def run():

    start = timer()

    log_list = get_eb_worker_logs('e-ev8fq2dhbv')
    for timestamp, url, instance_id in log_list:
        print("timestamp:", timestamp)
        print("url:", url)

        with urllib.request.urlopen(url) as download_file_obj:
            in_memory_file = io.BytesIO(download_file_obj.read())
            relevant_lines = process_zip_file(in_memory_file, instance_id)
            for i, line in enumerate(relevant_lines):
                print("\t", i, ":", line)

    end = timer()

    execution_seconds = end - start
    print("Completed in", execution_seconds, "seconds")


if __name__ == "__main__":

    run()
