import logging
from functools import wraps

import pyarrow.flight as fl

from feast.errors import FeastError

logger = logging.getLogger(__name__)


def arrow_client_error_handling_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            mapped_error = FeastError.from_error_detail(_get_exception_data(e.args[0]))
            if mapped_error is not None:
                raise mapped_error
            raise e

    return wrapper


def arrow_server_error_handling_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            if isinstance(e, FeastError):
                raise fl.FlightError(e.to_error_detail())
            # Re-raise non-Feast exceptions so Arrow Flight returns a proper error
            # instead of allowing the server method to return None.
            raise e

    return wrapper


def _get_exception_data(except_str) -> str:
    substring = "Flight error: "

    # Find the starting index of the substring
    position = except_str.find(substring)
    end_json_index = except_str.find("}")

    if position != -1 and end_json_index != -1:
        # Extract the part of the string after the substring
        result = except_str[position + len(substring) : end_json_index + 1]
        return result

    return ""