# coding: utf-8

__author__ = "Ondrej Jurcak"

import logging
from configparser import ConfigParser

from aws_xray_sdk import global_sdk_config
from aws_xray_sdk.core.async_context import AsyncContext
from aws_xray_sdk.core import xray_recorder, patch_all, patch
from aws_xray_sdk.core.utils import stacktrace
from aws_xray_sdk.ext.flask.middleware import XRayMiddleware
from aws_xray_sdk.ext.flask_sqlalchemy.query import XRayFlaskSqlAlchemy
from aws_xray_sdk.ext.util import construct_xray_header, inject_trace_header
from sentry_sdk.integrations.logging import ignore_logger

try:
    from celery import signals

    CELERY_INSTALLED = True
except ImportError:
    CELERY_INSTALLED = False


class XRaySettingsConfigParser:
    def __init__(self, settings):
        if "XRAY" in settings:
            xray = settings["XRAY"]
            self.enabled = xray.getboolean("ENABLED", False)
            self.service_name = xray.get("SERVICE_NAME")
            self.plugins = xray.get("PLUGINS", "'ECSPlugin', 'EC2Plugin'")
            self.sampling = xray.getboolean("SAMPLING", False)
            self.stream_sql = xray.getboolean("STREAM_SQL", True)
            self.daemon = xray.get("DAEMON", "xray-daemon:2000")
            self.patch_modules = xray.get("PATCH_MODULES", None)
            self.logging_enabled = xray.getboolean("LOGGING_ENABLED", False)
            self.sqlalchemy = xray.getboolean("SQLALCHEMY", False)
            self.middleware = xray.getboolean("MIDDLEWARE", True)
            self.sentry_logging = xray.getboolean("SENTRY_LOGGING", False)
            self.async_context = xray.getboolean("ASYNC_CONTEXT", False)
        else:
            self.enabled = False
            self.service_name = None
            self.plugins = "'ECSPlugin', 'EC2Plugin'"
            self.sampling = False
            self.stream_sql = True
            self.daemon = "xray-daemon:2000"
            self.patch_modules = None
            self.logging_enabled = False
            self.sqlalchemy = False
            self.middleware = True
            self.sentry_logging = False
            self.async_context = False


class XRaySettingsLazySettings:
    def __init__(self, settings):
        self.enabled = settings.XRAY_ENABLED
        self.service_name = settings.XRAY_SERVICE_NAME
        self.plugins = settings.XRAY_PLUGINS
        self.sampling = settings.XRAY_SAMPLING
        self.stream_sql = settings.XRAY_STREAM_SQL
        self.daemon = settings.XRAY_DAEMON
        self.patch_modules = settings.XRAY_PATCH_MODULES
        self.logging_enabled = settings.XRAY_LOGGING_ENABLED
        self.sqlalchemy = settings.XRAY_SQLALCHEMY
        self.middleware = settings.XRAY_MIDDLEWARE
        self.sentry_logging = settings.XRAY_SENTRY_LOGGING
        self.async_context = settings.XRAY_ASYNC_CONTEXT


class XRayConfigurator:
    def __init__(self, settings, application=None) -> None:
        logging.getLogger("xray_configurator_logger")
        if isinstance(settings, ConfigParser):
            self.settings = XRaySettingsConfigParser(settings)
        else:
            self.settings = XRaySettingsLazySettings(settings)

        if not self.settings.enabled:
            logging.info("xray sdk disabled")
            global_sdk_config.set_sdk_enabled(False)
            return

        self._init_xray(self.settings)

        if application:
            if self.settings.sqlalchemy:
                logging.info("XRayFlaskSqlAlchemy enabled")
                XRayFlaskSqlAlchemy(application)
            if self.settings.middleware:
                XRayMiddleware(application, xray_recorder)

    def _init_xray(self, settings):
        logging.info("xray sdk enabled " + str(settings.enabled))
        global_sdk_config.set_sdk_enabled(True)
        xray_recorder.configure(
            service=settings.service_name,
            plugins=eval(settings.plugins),
            sampling=settings.sampling,
            stream_sql=settings.stream_sql,
            daemon_address=settings.daemon,
            context_missing="LOG_ERROR",
            context=AsyncContext() if settings.async_context else None,
        )
        if not settings.patch_modules:
            patch_all()
        else:
            patch(eval(settings.patch_modules))

        if settings.logging_enabled:
            logging.info("XRay logging enabled")
            logging.getLogger("aws_xray_sdk").setLevel(logging.DEBUG)
        else:
            logging.info("XRay logging disabled")
            logging.getLogger("aws_xray_sdk").setLevel(logging.NOTSET)

        if not settings.sentry_logging:
            ignore_logger("aws_xray_sdk")
            ignore_logger("aws_xray_sdk.core.emitters.udp_emitter")
            ignore_logger("aws_xray_sdk.core.sampling.rule_poller")
            ignore_logger("aws_xray_sdk.core.sampling.target_poller")


if CELERY_INSTALLED:

    @signals.task_prerun.connect
    def task_prerun(task_id, task, *args, **kwargs):
        logging.info("task prerun task_id: " + task_id)
        xray_header = construct_xray_header(task.request)
        logging.info(
            "xray header constructed. Root: {}, Parent: {}".format(
                xray_header.root, xray_header.parent
            )
        )
        segment = xray_recorder.begin_segment(
            name=task.name,
            traceid=xray_header.root,
            parent_id=xray_header.parent,
        )
        segment.save_origin_trace_header(xray_header)
        segment.put_metadata("task_id", task_id, namespace="celery")

    @signals.task_postrun.connect()
    def task_postrun(task_id, *args, **kwargs):
        logging.info("task post run task_id: " + task_id)
        xray_recorder.end_segment()

    @signals.before_task_publish.connect
    def before_task_publish(sender, headers, **kwargs):
        logging.info("before taskpublish sender: " + sender)
        subsegment = xray_recorder.begin_subsegment(
            name=sender,
            namespace="remote",
        )

        if subsegment is None:
            logging.info("task taskpublish not in segment")
            return
        subsegment.put_metadata("task_id", headers.get("id"), namespace="celery")
        inject_trace_header(headers, subsegment)

    @signals.after_task_publish.connect
    def xray_after_task_publish(**kwargs):
        logging.info("after task publish")
        xray_recorder.end_subsegment()

    @signals.task_failure.connect
    def xray_task_failure(einfo, **kwargs):
        logging.info("task_failure")
        segment = xray_recorder.current_segment()
        if segment is None:
            logging.info("task failure not in segment")
            return
        if einfo:
            stack = stacktrace.get_stacktrace(limit=xray_recorder.max_trace_back)
            try:
                segment.add_exception(einfo.exception, stack)
            except FileNotFoundError as e:  # os.getcwd() crashes sometimes we need to handle this issue
                logging.info("File not found error whne segment.add_exception")
