import logging
import time
from functools import partial
from typing import Any, Dict, List, Optional, TypedDict

from starlette import status
from starlette.datastructures import Headers
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from owslib.dictutil import exclude
from owslib.logger.utils import get_status_code_log_level
from owslib.urlpath import is_path_match

logger = logging.getLogger(__name__)

safe_headers = partial(exclude, keys=["authorization"])


class RequestInfo(TypedDict, total=False):
    response: Message
    start_time: float
    end_time: float


class RequestLoggerMiddleware:
    LOG_MESSAGE = "{status} - {verb} {resource}"

    def __init__(
        self,
        app: ASGIApp,
        exclude_paths: Optional[List[str]] = None,
        extra_headers: bool = False,
    ):
        self.app = app
        self.exclude_paths = exclude_paths
        self.extra_headers = extra_headers

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or is_path_match(
            scope["path"], match=self.exclude_paths
        ):
            await self.app(scope, receive, send)
            return

        info = RequestInfo(response={})

        async def inner_send(message: Message) -> None:
            if message["type"] == "http.response.start":
                info["response"] = message
            await send(message)

        try:
            info["start_time"] = time.time()
            await self.app(scope, receive, inner_send)
        except Exception as exc:
            info["response"]["status"] = status.HTTP_500_INTERNAL_SERVER_ERROR
            raise exc
        finally:
            info["end_time"] = time.time()
            self.log(scope, info)

    def log(self, scope: Scope, info: RequestInfo) -> None:
        status_code = info["response"]["status"]
        level = get_status_code_log_level(status_code)
        extra: Dict[str, Any] = {
            "request.start_time": info["start_time"],
            "request.end_time": info["end_time"],
        }

        if self.extra_headers:
            filter_headers = partial(exclude, keys=["authorization"])
            request_headers = Headers(scope=scope)
            extra.update({"request.headers": filter_headers(dict(request_headers))})
            if response_headers := Headers(raw=info["response"].get("headers", [])):
                extra.update(
                    {"response.headers": filter_headers(dict(response_headers))}
                )

        logger.log(
            level,
            self.LOG_MESSAGE.format(
                status=status_code,
                verb=scope["method"],
                resource=scope["path"],
            ),
            extra=extra,
        )