# (c) Copyright IBM Corp. 2021
# (c) Copyright Instana Inc. 2020


try:
    import contextvars
    from typing import Any, Dict, Tuple
    from urllib import parse

    import celery  # noqa: F401
    from celery import registry, signals
    from opentelemetry import context, trace
    from opentelemetry.context import get_current

    from instana.log import logger
    from instana.propagators.format import Format
    from instana.singletons import get_tracer
    from instana.span.span import InstanaSpan
    from instana.util.traceutils import get_tracer_tuple

    client_token: Dict[str, Any] = {}
    worker_token: Dict[str, Any] = {}
    client_span = contextvars.ContextVar("client_span")
    worker_span = contextvars.ContextVar("worker_span")

    def _get_task_id(
        headers: Dict[str, Any],
        body: Tuple[str, Any],
    ) -> str:
        """
        Across Celery versions, the task id can exist in a couple of places.
        """
        id = headers.get("id", None)
        if id is None:
            id = body.get("id", None)
        return id

    def add_broker_attributes(
        span: InstanaSpan,
        broker_url: str,
    ) -> None:
        try:
            url = parse.urlparse(broker_url)

            # Add safety for edge case where scheme may not be a string
            url_scheme = str(url.scheme)
            span.set_attribute("scheme", url_scheme)

            span.set_attribute("host", url.hostname if url.hostname else "localhost")

            if not url.port:
                # Set default port if not specified
                if url_scheme == "redis":
                    span.set_attribute("port", "6379")
                elif "amqp" in url_scheme:
                    span.set_attribute("port", "5672")
                elif "sqs" in url_scheme:
                    span.set_attribute("port", "443")
            else:
                span.set_attribute("port", str(url.port))
        except Exception:
            logger.debug(f"Error parsing broker URL: {broker_url}", exc_info=True)

    @signals.task_prerun.connect
    def task_prerun(
        *args: Tuple[object, ...],
        **kwargs: Dict[str, Any],
    ) -> None:
        try:
            ctx = None
            tracer = get_tracer()

            task = kwargs.get("sender", None)
            task_id = kwargs.get("task_id", None)
            task = registry.tasks.get(task.name)

            headers = task.request.get("headers", {})
            if headers is not None:
                ctx = tracer.extract(
                    Format.HTTP_HEADERS, headers, disable_w3c_trace_context=True
                )

            span = tracer.start_span("celery-worker", context=ctx)
            span.set_attribute("task", task.name)
            span.set_attribute("task_id", task_id)
            add_broker_attributes(span, task.app.conf["broker_url"])

            ctx = trace.set_span_in_context(span)
            token = context.attach(ctx)
            worker_token["token"] = token
            worker_span.set(span)
        except Exception:
            logger.debug("celery-worker task_prerun: ", exc_info=True)

    @signals.task_postrun.connect
    def task_postrun(
        *args: Tuple[object, ...],
        **kwargs: Dict[str, Any],
    ) -> None:
        try:
            span = worker_span.get()

            if span.is_recording():
                span.end()
                worker_span.set(None)
            if "token" in worker_token:
                context.detach(worker_token.pop("token", None))
        except Exception:
            logger.debug("celery-worker after_task_publish: ", exc_info=True)

    @signals.task_failure.connect
    def task_failure(
        *args: Tuple[object, ...],
        **kwargs: Dict[str, Any],
    ) -> None:
        try:
            span = worker_span.get()
            if span.is_recording():
                span.set_attribute("success", False)
                exc = kwargs.get("exception", None)
                if exc:
                    span.record_exception(exc)
                else:
                    span.mark_as_errored()
        except Exception:
            logger.debug("celery-worker task_failure: ", exc_info=True)

    @signals.task_retry.connect
    def task_retry(
        *args: Tuple[object, ...],
        **kwargs: Dict[str, Any],
    ) -> None:
        try:
            span = worker_span.get()
            if span.is_recording():
                reason = kwargs.get("reason", None)
                if reason:
                    span.set_attribute("retry-reason", reason)
        except Exception:
            logger.debug("celery-worker task_failure: ", exc_info=True)

    @signals.before_task_publish.connect
    def before_task_publish(
        *args: Tuple[object, ...],
        **kwargs: Dict[str, Any],
    ) -> None:
        try:
            tracer, parent_span, _ = get_tracer_tuple()
            if not tracer:
                return

            parent_context = get_current()

            body = kwargs["body"]
            headers = kwargs["headers"]
            task_name = kwargs["sender"]
            task = registry.tasks.get(task_name)
            task_id = _get_task_id(headers, body)

            span = tracer.start_span("celery-client", context=parent_context)
            span.set_attribute("task", task_name)
            span.set_attribute("task_id", task_id)
            add_broker_attributes(span, task.app.conf["broker_url"])

            # Context propagation
            context_headers = {}
            tracer.inject(
                span.context,
                Format.HTTP_HEADERS,
                context_headers,
                disable_w3c_trace_context=True,
            )

            # Fix for broken header propagation
            # https://github.com/celery/celery/issues/4875
            task_headers = kwargs.get("headers") or {}
            task_headers.setdefault("headers", {})
            task_headers["headers"].update(context_headers)
            kwargs["headers"] = task_headers

            ctx = trace.set_span_in_context(span)
            token = context.attach(ctx)
            client_token["token"] = token
            client_span.set(span)
        except Exception:
            logger.debug("celery-client before_task_publish: ", exc_info=True)

    @signals.after_task_publish.connect
    def after_task_publish(
        *args: Tuple[object, ...],
        **kwargs: Dict[str, Any],
    ) -> None:
        try:
            span = client_span.get()
            if span.is_recording():
                span.end()
                client_span.set(None)
            if "token" in client_token:
                context.detach(client_token.pop("token", None))

        except Exception:
            logger.debug("celery-client after_task_publish: ", exc_info=True)

    logger.debug("Instrumenting celery")
except ImportError:
    pass
