"""Auth util."""
from typing import Callable
from functools import wraps

from flask import request
from jwtauth import JWTAuth
from jwtauth.utils import jwt_auth_from_environment
from jwtauth.exceptions import JWTAuthError
from oto import response
from oto.adaptors.flask import flaskify
from owsrequest import context as owsrequest_context

from product_digital import config


def only_for_identity(identity: str | list[str]) -> Callable:
    """Decorator function.

    Which checks if the identity of the user making the request
    matches the expected identities for the endpoint.

    .. code-block:: python
    @app.route("/")
    @only_for_identity("10436b38-5e11-472d-b6a4-bf1ee2b1b438")
    def index():
        return "Hello, Ben!"

    @app.route("/two")
    @only_for_identity([
        "10436b38-5e11-472d-b6a4-bf1ee2b1b438",
        "7f418dad-780b-4ab3-a4a2-2deba190f503"
    ])
    def index():
        return "Hello, Ben or Jess!"

    :param identity: The identity, or identities, of user(s) allowed to access the endpoint
    """
    identity_list: list[str] = [identity] if isinstance(identity, str) else identity

    def decorator(function: Callable) -> Callable:
        auth: JWTAuth = jwt_auth_from_environment(environment=config.ENVIRONMENT)

        @wraps(function)
        def wrapper(*args, **kwargs):
            try:
                context = owsrequest_context.get_request_context_from_headers(
                    request.headers, jwt_auth_client=auth
                )
                jwt_identity = context.jwt_identity_id
                if not jwt_identity:
                    raise JWTAuthError("No valid identity found in request context")

                if jwt_identity not in identity_list:
                    raise JWTAuthError(
                        f"Identity {jwt_identity} is not permitted for this endpoint"
                    )

            except JWTAuthError as auth_error:
                return flaskify(response.Response(message=auth_error.message, status=403))

            return function(*args, **kwargs)

        return wrapper

    return decorator