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

import time
import threading
from datetime import datetime, timedelta
import concurrent
import sys
import os
import asyncio
import logging

try:
    import psutil
except ImportError:
    pass  # If psutil isn't installed, simply does not capture process stats.

from azure.servicebus import ServiceBusMessage, ServiceBusMessageBatch
from azure.servicebus.exceptions import MessageAlreadySettled

from logger import get_logger
from app_insights_metric import AbstractMonitorMetric
from process_monitor import ProcessMonitor

PRINT_CONSOLE = True
LOGFILE_NAME = os.environ.get("DEBUG_SHARE") + "output"


class ReceiveType:
    push = "push"
    pull = "pull"
    none = None


class StressTestResults(object):
    def __init__(self):
        self.total_sent = 0
        self.total_received = 0
        self.time_elapsed = None
        self.state_by_sender = {}
        self.state_by_receiver = {}
        self.actual_size = 0

    def __repr__(self):
        return str(vars(self))


class StressTestRunnerState(object):
    """Per-runner state, e.g. if you spawn 3 senders each will have this as their state object,
    which will be coalesced at completion into StressTestResults"""

    def __init__(self):
        self.total_sent = 0
        self.total_received = 0
        self.cpu_percent = None
        self.memory_bytes = None
        self.timestamp = None
        self.exceptions = []

    def __repr__(self):
        return str(vars(self))

    def populate_process_stats(self, monitor):
        self.timestamp = datetime.utcnow()
        try:
            self.cpu_percent = psutil.cpu_percent()
            self.memory_bytes = psutil.virtual_memory().percent
        except NameError:
            return  # psutil was not installed, fall back to simply not capturing these stats.


class StressTestRunner:
    """Framework for running a service bus stress test.
    Duration can be overriden via the --stress_test_duration flag from the command line"""

    def __init__(
        self,
        senders,
        receivers,
        admin_client,
        duration=timedelta(minutes=15),
        receive_type=ReceiveType.push,
        send_batch_size=None,
        message_size=10,
        max_wait_time=10,
        send_delay=1.0,
        receive_delay=0,
        should_complete_messages=True,
        max_message_count=10,
        send_session_id=None,
        fail_on_exception=True,
        azure_monitor_metric=None,
        process_monitor=None,
        logging_level=logging.ERROR,
        transport_type=False,
        rotating_logs=True,
    ):
        self.senders = senders
        self.receivers = receivers
        self.admin_client = admin_client
        self.duration = duration
        self.receive_type = receive_type
        self.message_size = message_size
        self.send_batch_size = send_batch_size
        self.max_wait_time = max_wait_time
        self.send_delay = send_delay
        self.receive_delay = receive_delay
        self.should_complete_messages = should_complete_messages
        self.max_message_count = max_message_count
        self.fail_on_exception = fail_on_exception
        self.send_session_id = send_session_id
        self.azure_monitor_metric = azure_monitor_metric or AbstractMonitorMetric("fake_test_name")
        self.logging_level = logging_level
        self.rotating_logs = rotating_logs
        logfile_name = LOGFILE_NAME
        if transport_type:
            logfile_name += "_ws.log"
        else:
            logfile_name += ".log"
        self.logger = get_logger(logfile_name, "stress_test", self.logging_level, rotating_logs=self.rotating_logs)
        self.process_monitor = process_monitor or ProcessMonitor(
            "monitor_{}".format(logfile_name),
            "stress_test_queues",
            print_console=PRINT_CONSOLE,
        )

        # Because of pickle we need to create a state object and not just pass around ourselves.
        # If we ever require multiple runs of this one after another, just make Run() reset this.
        self._state = StressTestRunnerState()

        self._duration_override = None
        for arg in sys.argv:
            if arg.startswith("--duration="):
                self._duration_override = timedelta(seconds=int(arg.split("=")[1]))

        self._should_stop = False

    # Plugin functions the caller can override to further tailor the test.
    def on_send(self, state, sent_message, sender):
        """Called on every successful send, per message"""
        pass

    def on_receive(self, state, received_message, receiver):
        """Called on every successful receive, per message"""
        pass

    def on_receive_batch(self, state, batch, receiver):
        """Called on every successful receive, at the batch or iterator level rather than per-message"""
        pass

    def post_receive(self, state, receiver):
        """Called after completion of every successful receive"""
        pass

    def on_complete(self, send_results=[], receive_results=[]):
        """Called on stress test run completion"""
        pass

    def pre_process_message(self, message):
        """Allows user to transform the message before batching or sending it."""
        pass

    def pre_process_message_batch(self, message):
        """Allows user to transform the batch before sending it."""
        pass

    def pre_process_message_body(self, payload):
        """Allows user to transform message payload before sending it."""
        return payload

    def _schedule_interval_logger(self, end_time, logger, process_monitor, description="", interval_seconds=300):
        def _do_interval_logging():
            if end_time > datetime.utcnow() and not self._should_stop:
                self._state.populate_process_stats(process_monitor)
                logger.critical("%s RECURRENT STATUS: %s", description, self._state)
                self._schedule_interval_logger(end_time, logger, process_monitor, description, interval_seconds)

        t = threading.Timer(interval_seconds, _do_interval_logging)
        t.start()

    def _construct_message(self):
        if self.send_batch_size != None:
            batch = ServiceBusMessageBatch()
            for _ in range(self.send_batch_size):
                message = ServiceBusMessage(self.pre_process_message_body("a" * self.message_size))
                self.pre_process_message(message)
                batch.add_message(message)
            self.pre_process_message_batch(batch)
            return batch
        else:
            message = ServiceBusMessage(self.pre_process_message_body("a" * self.message_size))
            self.pre_process_message(message)
            return message

    def _send(self, sender, end_time):
        self._schedule_interval_logger(end_time, self.logger, self.process_monitor, "Sender " + str(self))
        try:
            self.logger.debug("Starting send loop")
            # log sender
            self.logger.debug("Sender: %s", sender)
            with sender:
                while end_time > datetime.utcnow() and not self._should_stop:
                    try:
                        message = self._construct_message()
                        if self.send_session_id != None:
                            message.session_id = self.send_session_id
                        self.logger.debug("Sending message: %s", message)
                        sender.send_messages(message)
                        self.azure_monitor_metric.record_messages_cpu_memory(
                            self.send_batch_size,
                            self.process_monitor.cpu_usage_percent,
                            self.process_monitor.memory_usage_percent,
                        )
                        if self.send_batch_size:
                            self._state.total_sent += self.send_batch_size
                        else:
                            self._state.total_sent += 1  # send single message
                        self.on_send(self._state, message, sender)

                    except Exception as e:
                        self.logger.exception("Exception during send: %s", e)
                        self.azure_monitor_metric.record_error(e)
                        self._state.exceptions.append(e)
                        if self.fail_on_exception:
                            raise
                    time.sleep(self.send_delay)
            self._state.timestamp = datetime.utcnow()
            return self._state
        except Exception as e:
            self.logger.exception("Exception in sender: %s", e)
            self._should_stop = True
            raise

    def _receive(self, receiver, end_time):
        self._schedule_interval_logger(end_time, self.logger, self.process_monitor, "Receiver " + str(self))
        # log receiver
        self.logger.debug("Receiver: %s", receiver)
        try:
            with receiver:
                while end_time > datetime.utcnow() and not self._should_stop:
                    try:
                        if self.receive_type == ReceiveType.pull:
                            batch = receiver.receive_messages(
                                max_message_count=self.max_message_count,
                                max_wait_time=self.max_wait_time,
                            )
                        elif self.receive_type == ReceiveType.push:
                            receiver.max_wait_time = self.max_wait_time
                            batch = receiver
                        # else:
                        #     batch = []

                        for message in batch:
                            # log reciever
                            self.logger.debug("Received message: %s", message)
                            self.on_receive(self._state, message, receiver)
                            try:
                                if self.should_complete_messages:
                                    receiver.complete_message(message)
                            except MessageAlreadySettled:  # It may have been settled in the plugin callback.
                                pass

                            self._state.total_received += 1
                            # TODO: Get EnqueuedTimeUtc out of broker properties and calculate latency. Should properties/app properties be mostly None?
                            if end_time <= datetime.utcnow():
                                break
                            time.sleep(self.receive_delay)
                            self.azure_monitor_metric.record_messages_cpu_memory(
                                1,
                                self.process_monitor.cpu_usage_percent,
                                self.process_monitor.memory_usage_percent,
                            )
                        self.post_receive(self._state, receiver)
                    except Exception as e:
                        self.logger.exception("Exception during receive: %s", e)
                        self._state.exceptions.append(e)
                        self.azure_monitor_metric.record_error(e)
                        if self.fail_on_exception:
                            raise
                self._state.timestamp = datetime.utcnow()
            return self._state
        except Exception as e:
            self.azure_monitor_metric.record_error(e)
            self.logger.exception("Exception in receiver %s", e)
            self._should_stop = True
            raise

    def run(self):

        start_time = datetime.utcnow()
        if isinstance(self.duration, int):
            self.duration = timedelta(seconds=self.duration)
        end_time = start_time + (self._duration_override or self.duration)

        with self.process_monitor:
            with concurrent.futures.ThreadPoolExecutor(max_workers=4) as proc_pool:
                self.logger.info("STARTING PROC POOL")
                if self.senders:
                    senders = [proc_pool.submit(self._send, sender, end_time) for sender in self.senders]
                else:
                    senders = []
                if self.receivers:
                    receivers = [proc_pool.submit(self._receive, receiver, end_time) for receiver in self.receivers]
                else:
                    receivers = []
                result = StressTestResults()
                for each in concurrent.futures.as_completed(senders + receivers):
                    self.logger.info("SOMETHING FINISHED")
                    if each in senders:
                        result.state_by_sender[each] = each.result()
                    if each in receivers:
                        result.state_by_receiver[each] = each.result()
                # TODO: do as_completed in one batch to provide a way to short-circuit on failure.
                if self.senders:
                    result.state_by_sender = {
                        s: f.result() for s, f in zip(self.senders, concurrent.futures.as_completed(senders))
                    }
                    result.total_sent = sum([r.total_sent for r in result.state_by_sender.values()])
                if self.receivers:
                    result.state_by_receiver = {
                        r: f.result() for r, f in zip(self.receivers, concurrent.futures.as_completed(receivers))
                    }
                    self.logger.info("Got receiver results")
                    result.total_received = sum([r.total_received for r in result.state_by_receiver.values()])
                result.time_elapsed = end_time - start_time
                self.logger.critical("Stress test completed.  Results:\n %s", result)
                return result


class StressTestRunnerAsync(StressTestRunner):
    def __init__(
        self,
        senders,
        receivers,
        duration=timedelta(minutes=15),
        admin_client=None,
        receive_type=ReceiveType.push,
        send_batch_size=None,
        message_size=10,
        max_wait_time=10,
        send_delay=1.00,
        receive_delay=0,
        should_complete_messages=True,
        max_message_count=1,
        send_session_id=None,
        fail_on_exception=True,
        azure_monitor_metric=None,
        process_monitor=None,
        logging_level=logging.ERROR,
        transport_type=False,
        rotating_logs=True,
    ):
        super(StressTestRunnerAsync, self).__init__(
            senders,
            receivers,
            duration=duration,
            admin_client=admin_client,
            receive_type=receive_type,
            send_batch_size=send_batch_size,
            message_size=message_size,
            max_wait_time=max_wait_time,
            send_delay=send_delay,
            receive_delay=receive_delay,
            should_complete_messages=should_complete_messages,
            max_message_count=max_message_count,
            send_session_id=send_session_id,
            fail_on_exception=fail_on_exception,
            azure_monitor_metric=azure_monitor_metric,
            process_monitor=process_monitor,
            logging_level=logging_level,
            transport_type=transport_type,
            rotating_logs=rotating_logs,
        )

    async def _send_async(self, sender, end_time):
        self._schedule_interval_logger(end_time, self.logger, self.process_monitor, "Sender " + str(self))
        try:
            async with sender:
                while end_time > datetime.utcnow() and not self._should_stop:
                    try:
                        message = self._construct_message()
                        if self.send_session_id != None:
                            message.session_id = self.send_session_id
                        await sender.send_messages(message)
                        self.azure_monitor_metric.record_messages_cpu_memory(
                            self.send_batch_size,
                            self.process_monitor.cpu_usage_percent,
                            self.process_monitor.memory_usage_percent,
                        )
                        if self.send_batch_size:
                            self._state.total_sent += self.send_batch_size
                        else:
                            self._state.total_sent += 1
                        self.on_send(self._state, message, sender)
                    except Exception as e:
                        self.logger.exception("Exception during send: %s", e)
                        self.azure_monitor_metric.record_error(e)
                        self._state.exceptions.append(e)
                        if self.fail_on_exception:
                            raise
                    await asyncio.sleep(self.send_delay)
            self._state.timestamp = datetime.utcnow()
            return self._state
        except Exception as e:
            self.logger.exception("Exception in sender: %s", e)
            self._should_stop = True
            raise

    async def _receive_handle_message(self, message, receiver, end_time):
        self.on_receive(self._state, message, receiver)
        try:
            if self.should_complete_messages:
                await receiver.complete_message(message)
        except MessageAlreadySettled:  # It may have been settled in the plugin callback.
            pass
        self._state.total_received += 1
        # TODO: Get EnqueuedTimeUtc out of broker properties and calculate latency. Should properties/app properties be mostly None?
        await asyncio.sleep(self.receive_delay)
        self.azure_monitor_metric.record_messages_cpu_memory(
            1,
            self.process_monitor.cpu_usage_percent,
            self.process_monitor.memory_usage_percent,
        )

    async def _receive_async(self, receiver, end_time):
        self._schedule_interval_logger(end_time, self.logger, self.process_monitor, "Receiver " + str(self))
        try:
            async with receiver:
                while end_time > datetime.utcnow() and not self._should_stop:
                    try:
                        if self.receive_type == ReceiveType.pull:
                            batch = await receiver.receive_messages(
                                max_message_count=self.max_message_count,
                                max_wait_time=self.max_wait_time,
                            )
                            for message in batch:
                                await self._receive_handle_message(message, receiver, end_time)
                        elif self.receive_type == ReceiveType.push:
                            receiver.max_wait_time = self.max_wait_time
                            batch = receiver
                            async for message in batch:
                                if end_time <= datetime.utcnow():
                                    break
                                await self._receive_handle_message(message, receiver, end_time)
                        self.post_receive(self._state, receiver)
                    except Exception as e:
                        self.logger.exception("Exception during receive: %s", e)
                        self._state.exceptions.append(e)
                        self.azure_monitor_metric.record_error(e)
                        if self.fail_on_exception:
                            raise
            self._state.timestamp = datetime.utcnow()
            return self._state
        except Exception as e:
            self.azure_monitor_metric.record_error(e)
            self.logger.exception("Exception in receiver %s", e)
            self._should_stop = True
            raise

    async def run_async(self):
        start_time = datetime.utcnow()
        if isinstance(self.duration, int):
            self.duration = timedelta(seconds=self.duration)
        end_time = start_time + (self._duration_override or self.duration)
        if self.senders:
            send_tasks = [asyncio.create_task(self._send_async(sender, end_time)) for sender in self.senders]
        else:
            send_tasks = []
        if self.receivers:
            receive_tasks = [
                asyncio.create_task(self._receive_async(receiver, end_time)) for receiver in self.receivers
            ]
        else:
            receive_tasks = []
        with self.process_monitor:
            # await asyncio.gather(*send_tasks, *receive_tasks)
            for task in asyncio.as_completed(send_tasks + receive_tasks):
                try:
                    await task
                except Exception as e:
                    print(e)
            result = StressTestResults()
            if self.senders:
                result.state_by_sender = {s: f.result() for s, f in zip(self.senders, send_tasks)}
                result.total_sent = sum([r.total_sent for r in result.state_by_sender.values()])
            if self.receivers:
                result.state_by_receiver = {r: f.result() for r, f in zip(self.receivers, receive_tasks)}
                self.logger.info("got receiver results")

                result.total_received = sum([r.total_received for r in result.state_by_receiver.values()])
            result.time_elapsed = end_time - start_time
            self.logger.critical("Stress test completed.  Results:\n%s", result)
            return result