import asyncio
import json
import time
from asyncio import AbstractEventLoop
from threading import Thread

import psutil
from aio_pika import Message
from aio_pika.abc import AbstractRobustChannel

from amqp.adapter.logging_utils import logging_info, logging_warning
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ScalingRequestAlert
from amqp.router.utils import await_future, await_result

CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 0.15, 0.90
RESOURCE_UNSET, RESOURCE_IDLE, RESOURCE_ACTIVE, RESOURCE_OVERLOAD = -1, 0, 1, 2


class ScaleRequestV1:

    def __init__(
        self,
        serviceId: str,
        taskId: str,
        maxAvailability: int,
        currentAvailability: int,
        requestType: ScalingRequestAlert,
    ):
        self.version = 1  # Version of the request, currently 1
        self.serviceId = serviceId
        self.taskId = taskId
        self.maxAvailability = maxAvailability
        self.currentAvailability = currentAvailability
        self.requestType = requestType
        self.thread = None

    def to_json(self) -> str:
        """Converts the object to a JSON string matching the Java structure"""
        return json.dumps(
            {
                "version": self.version,
                "serviceId": self.serviceId,
                "taskId": self.taskId,
                "maxAvailability": self.maxAvailability,
                "currentAvailability": self.currentAvailability,
                "requestType": self.requestType.value,  # Using .value for Enum serialization
            },
            indent=2,
        )


class RunningAverage:
    """
    A class to maintain a running average of the last N values.
    The intended use is to track CPU usage over time window, that is directly related to sample rate and window size.
    This is the reason why the constructor takes the number of items to store, calculated as window size / sample rate
    Initial_value is for unit test to emulate OVERLOAD or IDLE condition.
    """

    def __init__(self, num_items, initial_value=0):
        self.buffer = [initial_value] * num_items  # Fixed-size array
        self.pointer = 0  # Current position in array
        self.num_items = num_items  # Maximum number of items to store
        self.is_filled = False  # Track if buffer is fully filled

    def add(self, value):
        # Store the value at current pointer position
        self.buffer[self.pointer] = value

        # Move pointer to next position with wrap-around
        self.pointer += 1
        if self.pointer >= self.num_items:
            self.pointer = 0
            self.is_filled = True

    def get_average(self):
        # Determine how many items we should average
        count = self.num_items if self.is_filled else self.pointer

        if count == 0:
            return 0.0  # Avoid division by zero

        return sum(self.buffer) / count

    def get_current_values(self):
        """Returns the values in chronological order (oldest first)"""
        if not self.is_filled:
            return self.buffer[: self.pointer]

        return self.buffer[self.pointer :] + self.buffer[: self.pointer]


class BackpressureHandler:
    # Track the number of messages currently being processed
    current_parallel_executions = 0
    # helps detect IDLE condition
    last_data_message_time = 0
    # helps prevent flooding the system with backpressure events
    last_backpressure_event_time = 0
    last_backpressure_event = ScalingRequestAlert.UPDATE
    # _callback_list is used in unit tests to record the invoked callbacks
    _callback_list = None

    def __init__(
        self,
        channel: AbstractRobustChannel,
        loop: AbstractEventLoop,
        config: AMQConfiguration,
    ):
        self.lock = None
        self.channel = channel
        self.loop = loop
        self.config = config
        self.exchange = None
        self.swarm_service_id = self.config.amq_adapter.swarm_service_id
        self.swarm_task_id = self.config.amq_adapter.swarm_task_id
        self.parallel_workers = self.config.backpressure.threshold_threads
        self.average_cpu_usage = None
        self.do_loop = -1
        self._resource_usage_state = RESOURCE_UNSET
        self._resource_usage_changed = 0
        self._resource_average_value = 0
        self._last_resource_max_value = 100  # Default max value for CPU usage

    def _do_loop(self) -> bool:
        """
        Check if the loop should continue running.
        Positive value of do_loop indicates the number of iterations left.
        Negative value indicates the loop should run indefinitely.
        """
        _val = self.do_loop != 0
        if self.do_loop > 0:
            self.do_loop -= 1
        return _val

    def increase_parallel_executions(self):
        """Increase the number of parallel executions"""
        logging_info(
            "Backpressure: Increase parallel executions, current=%s",
            self.current_parallel_executions,
        )
        self.current_parallel_executions += 1

    def decrease_parallel_executions(self):
        """Decrease the number of parallel executions"""
        logging_info(
            "Backpressure: Decrease parallel executions, current=%s",
            self.current_parallel_executions,
        )
        if self.current_parallel_executions > 0:
            self.current_parallel_executions -= 1

    def update_parallel_executions(self, count: int):
        """Update the number of parallel executions"""
        self.current_parallel_executions = count

    def update_last_data_message_time(self):
        logging_info("Backpressure: Update last data message time")
        """Update the last data message time"""
        self.last_data_message_time = time.time()

    def start_backpressure_monitor(self) -> Thread:
        # Start the Backpressure monitor loop
        self.thread = Thread(target=self.backpressure_monitor_loop)
        self.thread.daemon = True  # This makes it a daemon thread
        self.thread.start()
        return self.thread

    async def check_overload_condition(self):
        """Check if the current parallel executions exceed the limit"""
        self.update_last_data_message_time()
        logging_info(
            "Backpressure: Check overload condition, current=%s, max=%s",
            self.current_parallel_executions,
            self.parallel_workers,
        )
        if self.current_parallel_executions >= self.parallel_workers - 1:
            # Check if the last backpressure event was not an overload
            if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD:
                # Trigger the overload event
                await self.handle_backpressure_overload_event()
                self.last_backpressure_event_time = time.time()
                self.last_backpressure_event = ScalingRequestAlert.OVERLOAD

    async def _backpressure_monitor(self):
        """Monitor the backpressure conditions and trigger events accordingly"""
        _time_window_millis = self.config.backpressure.time_window
        _idle_duration_millis = self.config.backpressure.idle_duration
        _overload_duration_millis = self.config.backpressure.cpu_overload_duration
        _monitor_interval = 0.5  # Monitor every 500ms
        CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
        IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90
        if self.average_cpu_usage is None:
            self.average_cpu_usage = RunningAverage(
                int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval)
            )
        _cpu_usage_state = CPU_ACTIVE
        _cpu_usage_changed = time.time()
        while self._do_loop():
            _current_cpu_usage = psutil.cpu_percent(interval=None)
            self.average_cpu_usage.add(_current_cpu_usage)
            _average_cpu_usage = self.average_cpu_usage.get_average()
            _now = time.time()
            # logging_info(
            #    "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s",
            #    _current_cpu_usage,
            #    _average_cpu_usage,
            #    _cpu_usage_state,
            #    self.last_data_message_time,
            #    self.last_backpressure_event,
            #    self.last_backpressure_event_time,
            # )
            if _average_cpu_usage < IDLE_THRESHOLD:
                if _cpu_usage_state != CPU_IDLE:
                    _cpu_usage_changed = _now
                _cpu_usage_state = CPU_IDLE
            elif _average_cpu_usage > OVERLOAD_THRESHOLD:
                if _cpu_usage_state != CPU_OVERLOAD:
                    _cpu_usage_changed = _now
                _cpu_usage_state = CPU_OVERLOAD
            else:
                if _cpu_usage_state != CPU_ACTIVE:
                    _cpu_usage_changed = _now
                _cpu_usage_state = CPU_ACTIVE

            # Check if the current time exceeds the last backpressure event time
            if (
                _cpu_usage_state == CPU_OVERLOAD
                and _now - _cpu_usage_changed > _overload_duration_millis
                and (
                    _now - self.last_backpressure_event_time > _time_window_millis
                    or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
                )
            ):
                # Trigger the overload event
                await self.handle_backpressure_overload_event()
                self.last_backpressure_event_time = _now
                self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
            elif (
                _cpu_usage_state == CPU_IDLE
                and _now - _cpu_usage_changed > _idle_duration_millis
                and _now - self.last_data_message_time > _idle_duration_millis
                and (
                    _now - self.last_backpressure_event_time > _time_window_millis
                    or self.last_backpressure_event != ScalingRequestAlert.IDLE
                )
            ):
                # Trigger the idle event
                await self._handle_backpressure_idle_event()
                self.last_backpressure_event = ScalingRequestAlert.IDLE
                self.last_backpressure_event_time = _now
            else:
                # Trigger update event
                if _now - self.last_backpressure_event_time > _time_window_millis:
                    # Trigger the update event
                    _scaling_request: ScaleRequestV1 = ScaleRequestV1(
                        self.swarm_service_id,
                        self.swarm_task_id,
                        100,
                        _average_cpu_usage,
                        ScalingRequestAlert.UPDATE,
                    )
                    # Address the message to any adapter capable of supporting BACKPRESSURE request
                    await self.publish_backpressure_request(_scaling_request)
                    self.last_backpressure_event_time = _now
                    self.last_backpressure_event = ScalingRequestAlert.UPDATE
            await asyncio.sleep(_monitor_interval)

    def backpressure_monitor_loop(self):
        _loop = asyncio.new_event_loop()
        _loop.run_until_complete(self._backpressure_monitor())

    async def handle_backpressure_overload_event(self):
        logging_warning("Backpressure: Capacity close to depleted!")
        # Send an Overload event to the Management service
        scaling_request: ScaleRequestV1 = ScaleRequestV1(
            self.swarm_service_id,
            self.swarm_task_id,
            self.parallel_workers,
            self.parallel_workers - self.current_parallel_executions,
            ScalingRequestAlert.OVERLOAD,
        )
        # Address the message to any management service instance capable of supporting BACKPRESSURE request
        await self.publish_backpressure_request(scaling_request)
        # update / reset time-window so that the OVERLOAD is not sent too often
        self.last_data_message_time = time.time()

    async def _handle_backpressure_idle_event(self):
        logging_warning("Backpressure: Service is idle.")
        # Send an Idle event to the Management service
        scaling_request: ScaleRequestV1 = ScaleRequestV1(
            self.swarm_service_id,
            self.swarm_task_id,
            self.parallel_workers,
            self.parallel_workers - self.current_parallel_executions,
            ScalingRequestAlert.IDLE,
        )
        # Address the message to any adapter capable of supporting BACKPRESSURE request
        await self.publish_backpressure_request(scaling_request)
        # update / reset time-window so that the IDLE is not sent too often
        self.last_data_message_time = time.time()

    async def publish_backpressure_request(self, scaling_request: ScaleRequestV1):
        # Publish the backpressure request to the management service
        logging_info(
            f"Publishing backpressure for {scaling_request.serviceId} with request type {scaling_request.requestType}"
        )
        if not self.exchange:

            async def _wrap_rabbit_mq_api_init(channel):
                _exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
                return _exchange

            if BackpressureHandler._callback_list is not None:
                BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
            self.exchange = await await_result(
                asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
            )

        if self.exchange:

            async def _wrap_rabbit_mq_api():
                if not self.channel.is_closed:
                    binary_content: bytes = scaling_request.to_json().encode("utf-8")
                    pika_message: Message = Message(
                        body=binary_content,
                        content_encoding="utf-8",
                        delivery_mode=2,
                        content_type="application/octet-stream",
                        headers=None,
                        priority=0,
                        correlation_id=None,
                    )
                    await self.exchange.publish(
                        message=pika_message, routing_key="backpressure-scaling-v1"
                    )
                    logging_info(
                        "Service Message Published to %s, msg: %s",
                        self.exchange.name,
                        str(binary_content),
                    )
                    return True
                return False

            if BackpressureHandler._callback_list is not None:
                BackpressureHandler._callback_list["_wrap_rabbit_mq_api"] = 1
            await await_future(asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop))
