import logging
from typing import Optional

from asgiref.local import Local
from asgiref.sync import iscoroutinefunction
from django.conf import settings
from django.utils import timezone
from django.utils.decorators import sync_and_async_middleware
from pythonjsonlogger import jsonlogger

# Create a local context storage to store trace ID for each request. This is
# local to a specific coroutine/thread.
local = Local()


def _get_trace_header(request) -> Optional[str]:
    return request.META.get("HTTP_X_CLOUD_TRACE_CONTEXT", None)


@sync_and_async_middleware
def TraceIdMiddleware(get_response):
    """Saves the Google Cloud trace context for use in logging."""
    if iscoroutinefunction(get_response):
        # async middlware to save trace ID
        async def async_middleware(request):
            """Extract trace_id and save it in local context"""
            local.trace_header = _get_trace_header(request)
            response = await get_response(request)
            return response

        return async_middleware

    # sync middleware that does the same thing
    def middleware(request):
        local.trace_header = _get_trace_header(request)
        response = get_response(request)
        return response

    return middleware


class TraceIdLoggingFilter(logging.Filter):
    """Add trace_id (from TraceIdMiddleware) to the logging context."""

    def filter(self, record: logging.LogRecord) -> bool:
        trace_header = getattr(local, "trace_header", None)
        if trace_header is not None:
            project = settings.GCP_PROJECT_NAME
            trace_key = trace_header.split("/")[0]
            record._trace_id = f"projects/{project}/traces/{trace_key}"
        else:
            record._trace_id = None
        return True


class CustomJsonFormatter(jsonlogger.JsonFormatter):
    """Format log records as JSON, so that Cloud Run's logging can parse
    them."""

    def add_fields(self, log_record, record, message_dict):
        super().add_fields(log_record, record, message_dict)

        # add the trace ID
        trace_id = getattr(record, "_trace_id", None)
        if trace_id is not None:
            log_record["logging.googleapis.com/trace"] = trace_id

        # the rest of this is copied from the PythonJSONLogger docs
        if not log_record.get("timestamp"):
            now = timezone.now()
            log_record["timestamp"] = now
        if log_record.get("level"):
            log_record["severity"] = log_record["level"].upper()
        else:
            log_record["severity"] = record.levelname
