from flask import Flask, Response, g, request

from owslib import constants, context


class OwsContext:
    def __init__(
        self,
        correlation_id_header: str = constants.HEADER_CORRELATION_ID,
        add_request_context: bool = True,
        label_profile: bool = False,
    ) -> None:
        self.correlation_id_header = correlation_id_header
        self.add_request_context = add_request_context
        self.label_profile = label_profile
        self._globals = g

    def init_app(self, app: Flask) -> None:
        app.before_request(self.set_correlation_id)
        app.after_request(self.add_correlation_id_to_response)

        if self.add_request_context:
            app.before_request(self.set_request_context)
            app.after_request(self.reset_request_context)

    def set_correlation_id(self) -> None:
        correlation_id = (
            request.headers.get(self.correlation_id_header)
            or context.correlation_id.get()
            or context.generate_correlation_id()
        )
        self._globals.correlation_id_token = context.correlation_id.set(correlation_id)

    def add_correlation_id_to_response(self, response: Response) -> Response:
        if correlation_id := context.correlation_id.get():
            response.headers[self.correlation_id_header] = correlation_id
        if "correlation_id_token" in self._globals:
            context.correlation_id.reset(self._globals.correlation_id_token)
            del self._globals.correlation_id_token
        return response

    def set_request_context(self) -> None:
        request_context = context.request_context_from_headers(
            headers=request.headers, label_profile=self.label_profile
        )
        self._globals.request_context_token = context.request_context.set(
            request_context
        )

    def reset_request_context(self, response: Response) -> Response:
        if "request_context_token" in self._globals:
            context.request_context.reset(self._globals.request_context_token)
            del self._globals.request_context_token
        return response