import time
import uuid
from typing import Dict

from starlette.datastructures import MutableHeaders
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from structlog import get_logger
from structlog.contextvars import bind_contextvars, clear_contextvars, get_contextvars


class ContextVars:
    __original_context_vars: Dict

    def __init__(self, **kwargs):
        self.__original_context_vars = {}
        self.__context_vars = kwargs

    def __enter__(self):
        self.__original_context_vars = get_contextvars()
        bind_contextvars(**self.__context_vars)

    def __exit__(self, *_):
        clear_contextvars()
        bind_contextvars(**self.__original_context_vars)


class RequestIdMiddleware(BaseHTTPMiddleware):
    """
    Middleware that injects a request id if none is set by caller.
    """

    def __init__(self, *args, **kwds) -> None:
        super().__init__(*args, **kwds)
        self._logger = get_logger("middleware.request_id")

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> None:
        request_id = request.headers.get("x-request-id")
        if not request.headers.get("x-request-id"):
            request_id = str(uuid.uuid4())
            log_call = self._logger.debug(
                "request id generated by application because none was provided",
                reqId=request_id,
            )
            if log_call:
                await log_call
            mutable_headers = MutableHeaders(request._headers)
            mutable_headers["X-Request-Id"] = request_id
            request._headers = mutable_headers
            request.scope.update(headers=request.headers.raw)

        inject_tracing_parameters = {"reqId": request_id}

        try:
            with ContextVars(**inject_tracing_parameters):
                response = await call_next(request)
        except Exception as e:
            raise e

        response.headers["X-Request-Id"] = request_id

        return response


class RequestLoggerMiddleware(BaseHTTPMiddleware):
    """
    Middleware that logs request information to structlog.
    """

    def __init__(self, *args, **kwds) -> None:
        super().__init__(*args, **kwds)
        self._logger = get_logger("middleware.request_logger")

    def _extract_request_metadata(self, request: Request) -> Dict:
        return {
            "agent": request.headers.get("user-agent"),
            "ip": str(request.client.host),
            "method": str(request.method),
            "path": str(request.url.path),
            "reqId": request.headers.get("x-request-id"),
            "uri": str(request.url),
        }

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> None:
        request_metadata = {
            **self._extract_request_metadata(request=request),
        }

        response = None
        start_time = time.perf_counter()
        try:
            response = await call_next(request)
        except Exception as error:
            raise error
        finally:
            end_time = time.perf_counter()
            request_metadata["duration"] = end_time - start_time

            if response:
                request_metadata["status"] = response.status_code
            
            log_call = self._logger.info("http request", **request_metadata)
            if log_call:
                await log_call

        return response
