import logging
from functools import wraps

import requests

from feast import RepoConfig
from feast.errors import FeastError
from feast.permissions.client.http_auth_requests_wrapper import (
    get_http_auth_requests_session,
)

logger = logging.getLogger(__name__)


def rest_error_handling_decorator(func):
    @wraps(func)
    def wrapper(config: RepoConfig, *args, **kwargs):
        assert isinstance(config, RepoConfig)

        # Get a Session object
        with get_http_auth_requests_session(config.auth_config) as session:
            # Define a wrapper for session methods
            def method_wrapper(method_name):
                original_method = getattr(session, method_name)

                @wraps(original_method)
                def wrapped_method(*args, **kwargs):
                    logger.debug(
                        f"Calling {method_name} with args: {args}, kwargs: {kwargs}"
                    )
                    response = original_method(*args, **kwargs)
                    logger.debug(
                        f"{method_name} response status code: {response.status_code}"
                    )

                    try:
                        response.raise_for_status()
                    except requests.RequestException:
                        logger.debug(f"response.json() = {response.json()}")
                        mapped_error = FeastError.from_error_detail(response.json())
                        logger.debug(f"mapped_error = {str(mapped_error)}")
                        if mapped_error is not None:
                            raise mapped_error
                    return response

                return wrapped_method

            # Enhance session methods
            session.get = method_wrapper("get")  # type: ignore[method-assign]
            session.post = method_wrapper("post")  # type: ignore[method-assign]
            session.put = method_wrapper("put")  # type: ignore[method-assign]
            session.delete = method_wrapper("delete")  # type: ignore[method-assign]

            # Pass the enhanced session object to the decorated function
            return func(session, config, *args, **kwargs)

    return wrapper