import dataclasses
import logging
import threading as thread
import uuid
from typing import Any, ClassVar, Dict, Set, Tuple, cast, Optional

#from ddtrace.helpers import get_correlation_ids
from ddtrace.tracer import Tracer

from chattr_config.monitoring.utils import datadog_enabled


@dataclasses.dataclass
class Context:
    ASSIGNABLE_FIELDS: ClassVar[Set[str]] = {"brand_id", "branduser_id"}
    _context_map: ClassVar[Dict[int, Tuple[object, object]]] = {}

    version_id: str = None
    brand_id: int = None
    branduser_id: int = None
    trace_id: str = None  # Used by the APM for tracing purposes that would differ from our internal one.
    correlation_id: str = None
    request_id: str = None
    impersonator: Optional[str] = None

    def __post_init__(self):
        self.version_id = self.version_id or 'v1'
        self.correlation_id = self.correlation_id or str(uuid.uuid4())
        if not self.trace_id and datadog_enabled():
            self.tracer = Tracer()
            #self.trace_id, span_id = get_correlation_ids()
            self.trace_id = self.tracer.get_log_correlation_context().get('trace_id')
            self.span_id = self.tracer.get_log_correlation_context().get('span_id')
        self.request_id = self.request_id or str(uuid.uuid4())

    def to_dict(self):
        return self.__dict__

    @classmethod
    def acquire_from_dict(cls, context_dict: Dict[str, Any], set_context: bool = False):
        # request_id is not assigned because this is on the receiving side of
        # the context, and thus, we would want to create a different request_id
        if context_dict:
            context = cls(
                version_id=context_dict.get("version_id"),
                correlation_id=context_dict.get("correlation_id"),
                brand_id=context_dict.get("brand_id"),
                trace_id=context_dict.get("trace_id"),
                branduser_id=context_dict.get("branduser_id"),
                impersonator=context_dict.get('impersonator'),
            )
        else:
            logging.warning(f"Attempted to acquire Context from {context_dict} as dict. Creating new Context")
            context = cls()

        if set_context:
            cls._context_map[thread.get_ident()] = (None, context)

        return context

    @classmethod
    def acquire(cls, request=None, fields: Optional[Dict[str, str]] = None) -> "Context":
        fields = fields or {}
        entry: Tuple[object, Context] = cast(Tuple[object, Context], cls._context_map.get(thread.get_ident()))
        if not entry:
            entry = cls._make_entry(request)
        request, context = entry
        usable_keys = set(fields.keys()).intersection(cls.ASSIGNABLE_FIELDS)
        for key in usable_keys:
            setattr(context, key, fields[key])

        _add_jwt_claims(request, context)

        return context

    @classmethod
    def release(cls, request=None):
        entry = cls._context_map.get(thread.get_ident())
        if not entry:
            logging.warning("Context mismatch where the entry does not exist.")
            return

        del Context._context_map[thread.get_ident()]
        if request and hasattr(request, "context"):
            delattr(request, "context")

    @classmethod
    def _make_entry(cls, request=None):
        if request:
            ctx = _make_context(request)
            request.context = ctx
        else:
            ctx = Context()
        entry = cls._context_map[thread.get_ident()] = (request, ctx)
        return entry


def _make_context(request) -> Context:

    headers = request.headers
    # https://docs.datadoghq.com/synthetics/apm
    trace_id = headers.get("x-datadog-trace-id")
    # https://en.wikipedia.org/wiki/List_of_HTTP_header_fields
    # Generated by Chattr.io servers.
    correlation_id = headers.get("X-Correlation-ID")

    # Generated by clients but can be defaulted to a server rendered version.
    request_id = headers.get("X-Request-ID")
    return Context(trace_id=trace_id, correlation_id=correlation_id, request_id=request_id)


def _add_jwt_claims(request, context) -> None:
    if request and getattr(request, 'authed_user'):
        context.impersonator = context.impersonator or getattr(request.authed_user, 'impersonator')
        context.brand_id = context.brand_id or getattr(request.authed_user, 'brand_id')
