from typing import (
    Collection,
    Dict,
    Optional,
    Callable,
    Union,
    Type,
    Iterable,
)

from messagebus import (
    Event,
    Message,
    MessageBus,
)
from opentelemetry import trace
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.trace import (
    SpanContext,
    SpanKind,
    Span,
)
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

tracer = trace.get_tracer(__name__)


class MessageBusInstrumentor(BaseInstrumentor):

    def __init__(
        self,
        component_name: Optional[str] = None,
    ):
        self.component_name = component_name

    def instrumentation_dependencies(
        self
    ) -> Collection[str]:
        return ()

    def _instrument(
        self,
        to_wrap: Union[Type[MessageBus], Iterable[Type[MessageBus]]],
        **kwargs
    ):
        if hasattr(to_wrap, "__next__"):
            for cls_to_wrap in to_wrap:
                self._wrap_cls(cls_to_wrap=cls_to_wrap)
        else:
            self._wrap_cls(to_wrap)

    def _uninstrument(
        self,
        **kwargs
    ):
        if "to_wrap" in kwargs:
            for extra_class in kwargs["to_wrap"]:
                extra_class.publish_event = extra_class._original_publish_event  # noqa
                del extra_class._original_publish_event  # noqa
                extra_class.send_message = extra_class._original_send_message  # noqa
                del extra_class.send_message
                if hasattr(extra_class, "_open_telemetry_component_name"):
                    del extra_class._open_telemetry_component_name  # noqa

    def _wrap_cls(self, cls_to_wrap: Type[MessageBus]):
        if self.component_name is not None:
            cls_to_wrap._open_telemetry_component_name = self.component_name
        cls_to_wrap.publish_event = self._build_publish_event_wrapper(cls_to_wrap.publish_event)  # noqa
        cls_to_wrap.send_message = self._build_send_message_wrapper(cls_to_wrap.send_message)  # noqa

    def _build_publish_event_wrapper(
        self,
        original_publish_event: Callable[[MessageBus, Event], Message]
    ):
        def wrapper(
            message_bus: MessageBus,
            event: Event,
        ):
            return self._publish_event_wrapper(
                message_bus=message_bus,
                event=event,
                original_publish_event=original_publish_event,
            )

        return wrapper

    def _publish_event_wrapper(
        self,
        message_bus: MessageBus,
        event: Event,
        original_publish_event: Callable[[MessageBus, Event], Message]
    ):
        if hasattr(self, "_open_telemetry_component_name"):
            if event.synchronous:
                span_kind = SpanKind.CLIENT
            else:
                span_kind = SpanKind.PRODUCER

            with tracer.start_as_current_span(
                name=f"{self._open_telemetry_component_name}->{event.subject}",
                kind=span_kind,
            ) as span:  # type: Span
                span.set_attribute("subject", event.subject)
                span.set_attribute("type", event.event_type)

                return self._publish_event_with_context(
                    message_bus=message_bus,
                    event=event,
                    original_publish_event=original_publish_event,
                )
        else:
            return self._publish_event_with_context(
                message_bus=message_bus,
                event=event,
                original_publish_event=original_publish_event,
            )

    def _build_send_message_wrapper(
        self,
        original_send_message: Callable[[MessageBus, str, Message], Message],
    ):
        def wrapper(
            message_bus,
            component: str,
            message: Message
        ):
            return self._send_message_with_context(
                message_bus=message_bus,
                component=component,
                message=message,
                original_send_message=original_send_message,
            )

        return wrapper

    def _send_message_wrapper(
        self,
        message_bus: MessageBus,
        component: str,
        message: Message,
        original_send_message: Callable[[MessageBus, str, Message], Message],
    ) -> Message:
        if hasattr(self, "_open_telemetry_component_name"):
            with tracer.start_as_current_span(
                name=f"{self._open_telemetry_component_name}->{component}[{message.token}]",
                kind=SpanKind.CLIENT,
            ) as span:  # type: Span
                span.set_attribute("componentName", component)
                span.set_attribute("token", message.token)
                span.set_attribute("format", message.data_type)
                span.set_attribute("timeout", message.timeout)

                return self._send_message_with_context(
                    message_bus=message_bus,
                    component=component,
                    message=message,
                    original_send_message=original_send_message,
                )
        else:
            return self._send_message_with_context(
                message_bus=message_bus,
                component=component,
                message=message,
                original_send_message=original_send_message,
            )

    def _publish_event_with_context(
        self,
        message_bus: MessageBus,
        event: Event,
        original_publish_event: Callable[[MessageBus, Event], Message]
    ) -> Message:
        self._add_span_context_header(event.headers)
        return original_publish_event(message_bus, event)

    def _send_message_with_context(
        self,
        message_bus: MessageBus,
        component: str,
        message: Message,
        original_send_message: Callable[[MessageBus, str, Message], Message],
    ) -> Message:
        self._add_span_context_header(message.headers)
        return original_send_message(message_bus, component, message)

    def _add_span_context_header(
        self,
        headers: Dict[str, str]
    ):
        headers["spanContext"] = self._build_context_header()

    @staticmethod
    def _build_context_header():
        carrier = {}
        TraceContextTextMapPropagator().inject(carrier)
        return str(carrier)
