import asyncio
import os
import threading
import time
from asyncio import AbstractEventLoop, Future
from typing import Dict

from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange
from opentelemetry import trace
from opentelemetry.trace import Tracer, TraceState
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.serializer import map_as_string
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
    DataMessage,
    DataMessageMagicByte,
    DataResponse,
    MultipartDataMessage,
)
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient


def collect_trace_info():
    trace_map = {}
    TraceContextTextMapPropagator().inject(
        carrier=trace_map, context=trace.context_api.get_current()
    )
    return trace_map


def get_default_trace_state():
    return TraceState(entries={})


def set_text(carrier, key, value):
    carrier[key] = value


class DataMessageHandler:
    def __init__(
        self,
        service_adapter: CleverThisServiceAdapter,
        tracer: Tracer,
        message_factory: DataMessageFactory,
        service_message_factory: ServiceMessageFactory,
        reply_to_exchange: AbstractRobustExchange,
        rabbit_mq_client: RabbitMQClient,
        amq_configuration: AMQConfiguration,
        loop,
        backpressure_handler: BackpressureHandler,
        file_handler: FileHandler = StreamingFileHandler(),
    ):
        self.service_adapter: CleverThisServiceAdapter = service_adapter
        self.tracer = tracer
        self.message_factory: DataMessageFactory = message_factory
        self.service_message_factory: ServiceMessageFactory = service_message_factory
        self.reply_to_exchange: AbstractRobustExchange = reply_to_exchange
        self.rabbit_mq_client: RabbitMQClient = rabbit_mq_client
        self.rabbit_mq_url = amq_configuration.dispatch.rabbit_mq_url
        self.output_dir = amq_configuration.dispatch.download_dir
        self.loop = loop
        self.service_adapter.loop = loop
        # Dictionary to store outstanding Futures
        self.outstanding: Dict[str, asyncio.Future] = {}
        self.reply_to: str | None = None
        self.file_handler: FileHandler = file_handler
        self.backpressure_handler: BackpressureHandler = backpressure_handler

    async def inbound_data_message_callback_as_thread(self, message: AbstractIncomingMessage):
        """
        The callback method that is called when a DataMessage is received from RabbitMQ, starts new Thread.
        """
        threading.Thread(
            target=lambda: asyncio.run(self.inbound_data_message_callback_no_thread(message))
        ).start()

    async def inbound_data_message_callback_no_thread(self, message: AbstractIncomingMessage):
        """
        Receives the DataMessage as bytes from RabbitMQ queue and deserializes it to the DataMessage object
        Ensures Jaeger trace-id propagation.
        Invokes custom handler/mapper/adapter method that corresponds to this message
        :param message:
        :return:
        """
        start = time.time()
        amq_message = await self.reconstitute_data_message(message)
        if amq_message is not None:
            self.ensure_trace_info(amq_message)
            if amq_message.magic() == DataMessageMagicByte.HTTP_REQUEST.value:
                await self.run_http_request_magic(message, amq_message, self.loop)
            else:
                logging_error(
                    "Unknown MAGIC value: [%s], don't know how to interpret the message",
                    amq_message.magic(),
                )
                asyncio.run_coroutine_threadsafe(message.nack(requeue=False), self.loop)
                return
        else:
            logging_error("######[%s] No valid AMQ Message present.", message.delivery_tag)
        self.record_duration(start, message.delivery_tag)
        asyncio.run_coroutine_threadsafe(message.ack(multiple=False), self.loop)

    async def run_http_request_magic(
        self,
        message: AbstractIncomingMessage,
        amq_message: DataMessage,
        loop: AbstractEventLoop,
    ) -> None:
        """
        Synchronous version of on_message method. This is a wrapper around the async on_message method.
        It also counts parallel executions and handles backpressure.
        :param message: RabbitMQ raw message
        :param amq_message: DataMessage
        :param loop: Event loop
        :return: DataResponse
        """
        # Increment the counter for parallel executions and check resource availability
        self.backpressure_handler.increase_parallel_executions()
        await self.backpressure_handler.check_overload_condition()
        logging_info(
            f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]"
        )
        if message.reply_to:
            self.reply_to = message.reply_to

        try:
            _a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection)
            _response: DataResponse = await self.service_adapter.on_message(_a_message)
            logging_debug(f"Result in thread: {_response}")
            try:
                await self.send_reply(message.reply_to, _response)
            except Exception as e:
                logging_error(f"Processing message: {e}")
                asyncio.run_coroutine_threadsafe(message.nack(requeue=False), loop=loop)

            if isinstance(amq_message, MultipartDataMessage):
                # remove the downloaded files from the output_dir
                for file_id, filename in amq_message.extra_data.items():
                    try:
                        logging_info(
                            f"[{message.delivery_tag}] Deleting file: {file_id} in {filename}"
                        )
                        os.remove(filename)
                    except OSError as e:
                        logging_error(f"Deleting file {filename}: {e}")
        except Exception as e:
            logging_error(f"Request handler: {e}")

        self.backpressure_handler.decrease_parallel_executions()
        await self.backpressure_handler.check_overload_condition()

    async def reconstitute_data_message(
        self, message: AbstractIncomingMessage
    ) -> DataMessage | None:
        amq_message: DataMessage | None = None
        delivery_tag = message.delivery_tag
        addressing: int = message.headers.get("clevermicro.addressing", 0)
        if addressing == 0:
            amq_message = DataMessageFactory.from_bytes(message.body)
        else:
            amq_message = await DataMessageFactory.from_stream(
                message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.loop
            )
            logging_info(f"######[{delivery_tag}] AMQ Message from stream: {amq_message}")
            if amq_message is not None:
                extra_data = await self.file_handler.on_message(
                    message=message,
                    json_data=amq_message.body(),
                    rabbitmq_url=self.rabbit_mq_url,
                    loop=self.loop,
                    output_dir=self.output_dir,
                )
                logging_info(f"######[{delivery_tag}] Multipart download result: {extra_data}")
                amq_message = MultipartDataMessage(amq_message, extra_data)
        if amq_message:
            logging_info(
                "######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s",
                message.delivery_tag,
                message.consumer_tag,
                amq_message.id(),
                amq_message.method(),
                amq_message.path(),
                map_as_string(amq_message.trace_info()),
            )
        return amq_message

    def ensure_trace_info(self, amq_message: DataMessage):
        propagator = TraceContextTextMapPropagator()
        context = propagator.extract(
            carrier=amq_message.trace_info(), context=trace.context_api.get_current()
        )
        self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
        # logging_info("CTX=%s", context)
        # current_span = self.tracer.start_span("AMQ-Adapter.Processing.Request", parent=parent_span)
        # context_with_span = trace.set_span_in_context(current_span, context)
        # logging.info(" [*] CTX2=%s", context_with_span)
        # propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
        # logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s",
        #             map_as_string(amq_message.trace_info()), context_with_span, context)

    def record_duration(self, start, delivery_tag):
        global last_data_message_time
        last_data_message_time = time.time()
        logging_info(
            "######[%s] Done, single ACK: %sms.",
            delivery_tag,
            last_data_message_time - start,
        )

    async def send_reply(self, reply_to: str, response: DataResponse):
        """
        The Service side code. When the CleverXXX service has output ready, in DataResponse object,
        call this method to return the reply back to the caller
        :param reply_to:connection.channel
        :param response:
        :return:
        """
        if reply_to:
            pika_message: Message = Message(
                body=DataResponseFactory.serialize(response),
                correlation_id=response.id(),
                content_type=(
                    response.content_type()[0]
                    if isinstance(response.content_type(), list)
                    else response.content_type()
                ),
            )
            _res = asyncio.run_coroutine_threadsafe(
                self.reply_to_exchange.publish(message=pika_message, routing_key=reply_to),
                self.loop,
            )
            logging_debug(
                "###### DONE Reply Published to exchg:(%s) To=%s, res=%s",
                self.reply_to_exchange.name,
                reply_to,
                _res,
            )

    async def reply_received_callback(self, message: AbstractIncomingMessage):
        """
        The reply for DataMessage that we sent out earlier - must match the received response with
        the original request and respond to the caller.
        :param message:
        :return:
        """
        reply_id = message.correlation_id
        delivery_tag = message.delivery_tag
        debug = f"######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}"
        logging_debug(debug)
        # TODO - figure out implementation of this part, as this is sending reply back either to a service,
        #        or to the user
        logging_info(f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}")

        future: Future = self.outstanding.pop(reply_id, None)
        # if reply is None:
        # logging.warn(" [W] ######[%s] IGNORING REPLY. No outstanding request for ID %s",
        #             delivery_tag, reply_id)
        # else:
        # result = DataResponseFactory.from_bytes(body)
        # reply.tracing_span().end()
        # TODO - figure out the right response format (not needed for Clever Swarm v0.1)
        # that currently runs Java version, this method here is for the future compatibility to enable
        # internal, service-2-service communication, and the proper response type will be decided than
        # reply.response().set_result(Response.ok(result.body, content_type).build())
        if future:
            response: DataResponse = DataResponseFactory.from_bytes(message.body)
            # Complete the Future with the response data
            future.set_result(response.body())

        await message.ack(multiple=False)
