import ast
import logging

from messagebus import (
    Message,
    Event,
)
from messageprocessor import MessageProcessorCallback
from opentelemetry import trace
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

tracer = trace.get_tracer(__name__)


class OpenTelemetryProcessorCallback(MessageProcessorCallback):
    def __init__(self):
        self.span_context = None
        self.span = None

    def process_begun(
        self,
        event_name,
        request_id,
        message: Message = None,
        event: Event = None,
    ):
        self.span_context = tracer.start_as_current_span(
            name=event_name,
            context=self._get_parent_context(
                event=event,
                message=message
            ),
        )
        self.span = self.span_context.__enter__()

    def input_parsed(
        self,
        event_name,
        request_id,
        input_obj
    ):
        return

    def business_called(
        self,
        event_name,
        request_id,
        input_obj,
        output_obj
    ):
        return

    def process_finished(
        self,
        event_name,
        request_id,
        input_obj,
        output_obj
    ):
        self.span.set_status(status=StatusCode.OK)
        self.span_context.__exit__(None, None, None)

    def parse_exception(
        self,
        event_name,
        request_id,
        data,
        exception
    ):
        self._handle_exception(exception)

    def unhandled_exception(
        self,
        event_name,
        request_id,
        input_obj,
        exception
    ):
        self._handle_exception(exception)

    def _handle_exception(self, ex):
        self.span_context.__exit__(type(ex), ex, ex.__traceback__)

    @staticmethod
    def _get_parent_context(
        event,
        message
    ):
        headers = None
        if message is not None:
            headers = message.headers
        elif event is not None:
            headers = event.headers

        if headers is None or "spanContext" not in headers or headers["spanContext"] is None:
            return None

        try:
            carrier = ast.literal_eval(headers["spanContext"])
            return TraceContextTextMapPropagator().extract(carrier=carrier)
        except Exception as ex:
            logging.getLogger(__name__).exception(
                msg=f"Invalid spanContext: {headers['spanContext']}",
                exc_info=ex
            )
            return None
