# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

import os
import logging
import threading
import time
import asyncio
from argparse import ArgumentParser
from dotenv import load_dotenv

from azure.eventhub import EventHubProducerClient, EventData, EventHubSharedKeyCredential, TransportType
from azure.eventhub.exceptions import EventHubError
from azure.eventhub.aio import EventHubProducerClient as EventHubProducerClientAsync
from azure.identity import ClientSecretCredential, DefaultAzureCredential
from azure.identity.aio import (
    ClientSecretCredential as ClientSecretCredentialAsync,
    DefaultAzureCredential as DefaultAzureCredentialAsync,
)

from logger import get_logger
from process_monitor import ProcessMonitor
from app_insights_metric import AzureMonitorMetric

ENV_FILE = os.environ.get("ENV_FILE")


def handle_exception(error, ignore_send_failure, stress_logger, azure_monitor_metric):
    err_msg = "Sync send failed due to error: %s", repr(error)
    azure_monitor_metric.record_error(error)
    if ignore_send_failure:
        stress_logger.warning(err_msg)
        return 0
    raise error


def on_success(events, pid):
    # sending succeeded
    pass


def on_error(events, pid, error):
    # sending failed
    pass


async def on_success_async(events, pid):
    # sending succeeded
    pass


async def on_error_async(events, pid, error):
    # sending failed
    pass


def stress_send_sync(producer: EventHubProducerClient, args, stress_logger, azure_monitor_metric):
    try:
        batch = producer.create_batch(partition_id=args.send_partition_id, partition_key=args.send_partition_key)
        while True:
            event_data = EventData(body=b"D" * args.payload)
            batch.add(event_data)
    except ValueError:
        try:
            producer.send_batch(batch)
        except EventHubError as e:
            return handle_exception(e, args.ignore_send_failure, stress_logger, azure_monitor_metric)
    except EventHubError as e:
        return handle_exception(e, args.ignore_send_failure, stress_logger, azure_monitor_metric)
    return len(batch)


def stress_send_list_sync(producer: EventHubProducerClient, args, stress_logger, azure_monitor_metric):
    quantity = int(256 * 1023 / args.payload)
    send_list = []
    for _ in range(quantity):
        send_list.append(EventData(body=b"D" * args.payload))
    try:
        producer.send_batch(send_list)
    except EventHubError as e:
        return handle_exception(e, args.ignore_send_failure, stress_logger, azure_monitor_metric)
    return len(send_list)


async def stress_send_async(producer: EventHubProducerClientAsync, args, stress_logger, azure_monitor_metric):
    try:
        batch = await producer.create_batch(partition_id=args.send_partition_id, partition_key=args.send_partition_key)
        while True:
            event_data = EventData(body=b"D" * args.payload)
            batch.add(event_data)
    except ValueError:
        try:
            await producer.send_batch(batch)
        except EventHubError as e:
            return handle_exception(e, args.ignore_send_failure, stress_logger, azure_monitor_metric)
    except EventHubError as e:
        return handle_exception(e, args.ignore_send_failure, stress_logger, azure_monitor_metric)
    return len(batch)


async def stress_send_list_async(producer: EventHubProducerClientAsync, args, stress_logger, azure_monitor_metric):
    quantity = int(256 * 1023 / args.payload)
    send_list = []
    for _ in range(quantity):
        send_list.append(EventData(body=b"D" * args.payload))
    try:
        await producer.send_batch(send_list)
    except EventHubError as e:
        return handle_exception(e, args.ignore_send_failure, stress_logger, azure_monitor_metric)
    return len(send_list)


class StressTestRunner:
    def __init__(self, argument_parser):
        self.argument_parser = argument_parser
        self.argument_parser.add_argument("-m", "--method", default="stress_send_list_sync")
        self.argument_parser.add_argument(
            "--output_interval", type=float, default=int(os.environ.get("OUTPUT_INTERVAL", 5000))
        )
        self.argument_parser.add_argument(
            "--duration",
            help="Duration in seconds of the test",
            type=int,
            default=int(os.environ.get("DURATION", 999999999)),
        )
        self.argument_parser.add_argument(
            "--partitions", help="Number of partitions. 0 means to get partitions from eventhubs", type=int, default=0
        )
        self.argument_parser.add_argument("--send_partition_id", help="Send messages to a specific partition", type=int)
        self.argument_parser.add_argument(
            "--send_partition_key", help="Send messages to a specific partition key", type=str
        )
        self.argument_parser.add_argument(
            "--conn_str", help="EventHub connection string", default=os.environ.get("EVENT_HUB_CONN_STR")
        )
        self.argument_parser.add_argument("--azure_identity", help="Use identity", type=bool, default=False)
        parser.add_argument("--auth_timeout", help="Authorization Timeout", type=float, default=60)
        self.argument_parser.add_argument(
            "--eventhub", help="Name of EventHub", default=os.environ.get("EVENT_HUB_NAME")
        )
        self.argument_parser.add_argument(
            "--transport_type", help="Transport type, 0 means AMQP, 1 means AMQP over WebSocket", type=int, default=0
        )
        self.argument_parser.add_argument("--parallel_send_cnt", help="Number of parallelling sending", type=int)
        self.argument_parser.add_argument(
            "--parallel_create_new_client",
            action="store_true",
            help="Whether create new client for each sending",
        )
        self.argument_parser.add_argument("--buffered_mode", help="buffer producer", action="store_true")
        self.argument_parser.add_argument("--proxy_hostname", type=str)
        self.argument_parser.add_argument("--proxy_port", type=str)
        self.argument_parser.add_argument("--proxy_username", type=str)
        self.argument_parser.add_argument("--proxy_password", type=str)
        self.argument_parser.add_argument(
            "--hostname",
            help="The fully qualified host name for the Event Hubs namespace.",
            default=os.environ.get("EVENT_HUB_HOSTNAME"),
        )
        self.argument_parser.add_argument("--sas_policy", help="Name of the shared access policy to authenticate with")
        self.argument_parser.add_argument("--sas_key", help="Shared access key")
        self.argument_parser.add_argument("--aad_client_id", help="AAD client id")
        self.argument_parser.add_argument("--aad_secret", help="AAD secret")
        self.argument_parser.add_argument("--aad_tenant_id", help="AAD tenant id")
        self.argument_parser.add_argument("--payload", help="payload size", type=int, default=1024)
        self.argument_parser.add_argument("--pyamqp_logging_enable", help="pyamqp logging enable", action="store_true")
        self.argument_parser.add_argument("--print_console", action="store_true")
        self.argument_parser.add_argument("--log_filename", help="log file name", type=str)
        self.argument_parser.add_argument("--retry_total", type=int, default=3)
        self.argument_parser.add_argument("--retry_backoff_factor", type=float, default=0.8)
        self.argument_parser.add_argument("--retry_backoff_max", type=float, default=120)
        self.argument_parser.add_argument("--ignore_send_failure", help="ignore sending failures", action="store_true")
        self.argument_parser.add_argument("--uamqp_mode", help="Flag for uamqp or pyamqp", action="store_true")
        self.argument_parser.add_argument(
            "--debug_level",
            help="Flag for setting a debug level, can be Info, Debug, Warning, Error or Critical",
            type=str,
            default="Error",
        )
        # rotate logs by default, if you want to disable it, use --no-rotating-logs flag
        parser.add_argument("--no-rotating-logs", action="store_true")
        self.args, _ = parser.parse_known_args()

        # store rotating_logs in args for later use
        self.rotating_logs = not self.args.no_rotating_logs
        if self.args.send_partition_key and self.args.send_partition_id:
            raise ValueError("Cannot set send_partition_key and send_partition_id at the same time.")

        self.running = False

    def create_client(self, client_class, is_async=False):

        self.debug_level = getattr(logging, self.args.debug_level.upper(), logging.ERROR)
        transport_type = TransportType.Amqp if self.args.transport_type == 0 else TransportType.AmqpOverWebsocket
        http_proxy = None
        retry_options = {
            "retry_total": self.args.retry_total,
            "retry_backoff_factor": self.args.retry_backoff_factor,
            "retry_backoff_max": self.args.retry_backoff_max,
        }

        if self.args.proxy_hostname:
            http_proxy = {
                "proxy_hostname": self.args.proxy_hostname,
                "proxy_port": self.args.proxy_port,
                "username": self.args.proxy_username,
                "password": self.args.proxy_password,
            }
        if self.args.buffered_mode:
            if is_async:
                client = client_class(
                    fully_qualified_namespace=self.args.hostname,
                    eventhub_name=self.args.eventhub,
                    credential=DefaultAzureCredentialAsync(),
                    auth_timeout=self.args.auth_timeout,
                    http_proxy=http_proxy,
                    transport_type=transport_type,
                    logging_enable=self.args.pyamqp_logging_enable,
                    buffered_mode=self.args.buffered_mode,
                    on_success=on_success_async,
                    on_error=on_error_async,
                    uamqp_transport=self.args.uamqp_mode,
                    **retry_options,
                )
            else:
                client = client_class(
                    fully_qualified_namespace=self.args.hostname,
                    eventhub_name=self.args.eventhub,
                    credential=DefaultAzureCredential(),
                    auth_timeout=self.args.auth_timeout,
                    http_proxy=http_proxy,
                    transport_type=transport_type,
                    logging_enable=self.args.pyamqp_logging_enable,
                    buffered_mode=self.args.buffered_mode,
                    on_success=on_success,
                    on_error=on_error,
                    uamqp_transport=self.args.uamqp_mode,
                    **retry_options,
                )
        else:
            if is_async:
                client = client_class(
                    fully_qualified_namespace=self.args.hostname,
                    eventhub_name=self.args.eventhub,
                    credential=DefaultAzureCredentialAsync(),
                    auth_timeout=self.args.auth_timeout,
                    http_proxy=http_proxy,
                    transport_type=transport_type,
                    logging_enable=self.args.pyamqp_logging_enable,
                    uamqp_transport=self.args.uamqp_mode,
                    **retry_options,
                )
            else:
                client = client_class(
                    fully_qualified_namespace=self.args.hostname,
                    eventhub_name=self.args.eventhub,
                    credential=DefaultAzureCredential(),
                    auth_timeout=self.args.auth_timeout,
                    http_proxy=http_proxy,
                    transport_type=transport_type,
                    logging_enable=self.args.pyamqp_logging_enable,
                    uamqp_transport=self.args.uamqp_mode,
                    **retry_options,
                )

        return client

    def run(self):
        method_name = self.args.method
        if "async" in method_name:
            loop = asyncio.get_event_loop()
            loop.run_until_complete(self.run_async())
        else:
            self.run_sync()

    def run_sync(self):
        self.debug_level = getattr(logging, self.args.debug_level.upper(), logging.ERROR)
        log_filename = self.args.log_filename if self.args.log_filename else "producer_sync"
        if self.args.transport_type == 1:
            log_filename += "_ws.log"
        else:
            log_filename += ".log"
        with ProcessMonitor(
            "monitor_{}".format(log_filename), "producer_stress_sync", print_console=self.args.print_console
        ) as process_monitor:

            class EventHubProducerClientTest(EventHubProducerClient):
                def get_partition_ids(self_inner):
                    if self.args.partitions != 0:
                        return [str(i) for i in range(self.args.partitions)]
                    else:
                        return super(EventHubProducerClientTest, self_inner).get_partition_ids()

            method_name = self.args.method
            logdir = os.environ.get("DEBUG_SHARE")
            logfilepath = f"{logdir}/{log_filename}"

            logger = get_logger(
                logfilepath,
                method_name,
                level=self.debug_level,
                print_console=self.args.print_console,
                rotating_logs=self.rotating_logs,
            )
            test_method = globals()[method_name]
            self.running = True

            if self.args.parallel_send_cnt and self.args.parallel_send_cnt > 1:
                if self.args.parallel_create_new_client:
                    clients = [
                        self.create_client(EventHubProducerClientTest) for _ in range(self.args.parallel_send_cnt)
                    ]
                else:
                    clients = [self.create_client(EventHubProducerClientTest)]
                self.run_test_method_parallel(test_method, clients, logger, process_monitor)
            else:
                client = self.create_client(EventHubProducerClientTest)
                self.run_test_method(test_method, client, logger, process_monitor)

    def stop(self):
        self.running = False

    def run_test_method(self, test_method, worker, logger, process_monitor):
        deadline = time.time() + self.args.duration
        azure_monitor_metric = AzureMonitorMetric("Sync EventHubProducerClient")
        with worker:
            total_processed = 0
            iter_processed = 0
            start_time = time.perf_counter()
            while self.running and time.time() < deadline:
                try:
                    cur_iter_start_time = time.perf_counter()
                    processed = test_method(worker, self.args, logger, azure_monitor_metric)
                    now_time = time.perf_counter()
                    cur_iter_time_elapsed = now_time - cur_iter_start_time
                    total_processed += processed
                    iter_processed += processed
                    if iter_processed >= self.args.output_interval:
                        time_elapsed = now_time - start_time
                        logger.info(
                            "Total sent: %r, Time elapsed: %r, Speed: %r/s, Current Iteration Speed: %r/s",
                            total_processed,
                            time_elapsed,
                            total_processed / time_elapsed,
                            processed / cur_iter_time_elapsed,
                        )
                        azure_monitor_metric.record_events_cpu_memory(
                            iter_processed, process_monitor.cpu_usage_percent, process_monitor.memory_usage_percent
                        )
                        iter_processed -= self.args.output_interval
                except KeyboardInterrupt:
                    logger.info("keyboard interrupted")
                    self.stop()
                except Exception as e:
                    logger.exception("%r failed: %r", type(worker), e)
                    self.stop()
            logger.info("%r has finished testing", test_method)

    def run_test_method_parallel(self, test_method, workers, logger, process_monitor):
        cnt = self.args.parallel_send_cnt
        threads = []
        if self.args.parallel_create_new_client:
            for i in range(cnt):
                threads.append(
                    threading.Thread(
                        target=self.run_test_method,
                        args=(test_method, workers[i], logger, process_monitor),
                        daemon=True,
                    )
                )
        else:
            threads = [
                threading.Thread(
                    target=self.run_test_method, args=(test_method, workers[0], logger, process_monitor), daemon=True
                )
                for _ in range(cnt)
            ]

        for thread in threads:
            thread.start()

        for thread in threads:
            thread.join(timeout=self.args.duration)

    async def run_async(self):
        self.debug_level = getattr(logging, self.args.debug_level.upper(), logging.ERROR)
        log_filename = self.args.log_filename if self.args.log_filename else "producer_async"
        if self.args.transport_type == 1:
            log_filename += "_ws.log"
        else:
            log_filename += ".log"
        with ProcessMonitor(
            "monitor_{}".format(log_filename), "producer_stress_async", print_console=self.args.print_console
        ) as process_monitor:

            class EventHubProducerClientTestAsync(EventHubProducerClientAsync):
                async def get_partition_ids(self_inner):
                    if self.args.partitions != 0:
                        return [str(i) for i in range(self.args.partitions)]
                    else:
                        return await super(EventHubProducerClientTestAsync, self_inner).get_partition_ids()

            method_name = self.args.method
            logdir = os.environ.get("DEBUG_SHARE")
            logfilepath = f"{logdir}/{log_filename}"

            logger = get_logger(logfilepath, method_name, level=self.debug_level, print_console=self.args.print_console, rotating_logs=self.rotating_logs)
            test_method = globals()[method_name]
            self.running = True

            if self.args.parallel_send_cnt and self.args.parallel_send_cnt > 1:
                if self.args.parallel_create_new_client:
                    clients = [
                        self.create_client(EventHubProducerClientTestAsync, is_async=True)
                        for _ in range(self.args.parallel_send_cnt)
                    ]
                else:
                    clients = [self.create_client(EventHubProducerClientTestAsync, is_async=True)]
                await self.run_test_method_parallel_async(test_method, clients, logger, process_monitor)
            else:
                client = self.create_client(EventHubProducerClientTestAsync, is_async=True)
                await self.run_test_method_async(test_method, client, logger, process_monitor)

    async def run_test_method_async(self, test_method, worker, logger, process_monitor):
        deadline = time.time() + self.args.duration
        azure_monitor_metric = AzureMonitorMetric("Async EventHubProducerClient")
        async with worker:
            total_processed = 0
            iter_processed = 0
            start_time = time.perf_counter()
            while self.running and time.time() < deadline:
                try:
                    cur_iter_start_time = time.perf_counter()
                    processed = await test_method(worker, self.args, logger, azure_monitor_metric)
                    now_time = time.perf_counter()
                    cur_iter_time_elapsed = now_time - cur_iter_start_time
                    total_processed += processed
                    iter_processed += processed
                    if iter_processed >= self.args.output_interval:
                        time_elapsed = now_time - start_time
                        logger.info(
                            "Total sent: %r, Time elapsed: %r, Overall Speed: %r/s, Current Iteration Speed: %r/s",
                            total_processed,
                            time_elapsed,
                            total_processed / time_elapsed,
                            processed / cur_iter_time_elapsed,
                        )
                        azure_monitor_metric.record_events_cpu_memory(
                            iter_processed, process_monitor.cpu_usage_percent, process_monitor.memory_usage_percent
                        )
                        iter_processed -= self.args.output_interval
                except KeyboardInterrupt:
                    logger.info("keyboard interrupted")
                    self.stop()
                except Exception as e:
                    logger.exception("%r failed: %r", type(worker), e)
                    self.stop()
            logger.info("%r has finished testing", test_method)

    async def run_test_method_parallel_async(self, test_method, worker, logger, process_monitor):
        cnt = self.args.parallel_send_cnt
        tasks = []
        if self.args.parallel_create_new_client:
            for i in range(cnt):
                tasks.append(
                    asyncio.ensure_future(self.run_test_method_async(test_method, worker[i], logger, process_monitor))
                )
        else:
            tasks = [
                asyncio.ensure_future(self.run_test_method_async(test_method, worker[0], logger, process_monitor))
                for _ in range(cnt)
            ]

        await asyncio.gather(*tasks, return_exceptions=True)


if __name__ == "__main__":
    load_dotenv(dotenv_path=ENV_FILE, override=True)
    parser = ArgumentParser()
    runner = StressTestRunner(parser)
    runner.run()