from snowflake.connector.errors import (
    BadRequest,
    ForbiddenError,
    InternalServerError,
    RequestTimeoutError,
    TokenExpiredError,
    TooManyRequests,
)
from snowflake.connector.network import ReauthenticationRequest
from snowflake.core import exceptions as core_exceptions

__all__ = [
    'CortexSearchError',
    'CortexSearchConfigError',
    'CortexSearchAuthError',
    'CortexSearchTokenExpiredError',
    'CortexSearchBadRequestError',
    'CortexSearchRateLimitError',
    'CortexSearchServiceError',
    'CortexSearchTimeoutError',
    'CortexSearchNotFoundError',
]


class CortexSearchError(Exception):
    """Cortex Search base error."""


class CortexSearchConfigError(CortexSearchError):
    """Config error."""


class CortexSearchAuthError(CortexSearchError):
    """Authentication or authorization failed (401/403)."""


class CortexSearchTokenExpiredError(CortexSearchAuthError):
    """Session/JWT token expired — recoverable by re-authenticating.

    Treated as a transient error by the retry layer (see ``RETRYABLE_ERRORS``).
    Callers should generally catch the parent ``CortexSearchAuthError`` rather
    than this subclass directly.
    """


class CortexSearchBadRequestError(CortexSearchError):
    """Invalid query or filter (400)."""


class CortexSearchNotFoundError(CortexSearchError):
    """Upstream service error (404)."""


class CortexSearchRateLimitError(CortexSearchError):
    """Rate limit exceeded (429)."""


class CortexSearchServiceError(CortexSearchError):
    """Upstream service error (500)."""


class CortexSearchTimeoutError(CortexSearchError):
    """Request timed out."""


_SNOWFLAKE_ERROR_MAP = [
    # snowflake.connector errors
    (BadRequest, CortexSearchBadRequestError),
    (TokenExpiredError, CortexSearchTokenExpiredError),
    (ForbiddenError, CortexSearchAuthError),
    (TooManyRequests, CortexSearchRateLimitError),
    (InternalServerError, CortexSearchServiceError),
    (RequestTimeoutError, CortexSearchTimeoutError),
    # snowflake.core errors (raised by the REST/SDK layer);
    # 401 Unauthorized = expired session token (recoverable),
    # 403 Forbidden = permission denied (not recoverable).
    (core_exceptions.UnauthorizedError, CortexSearchTokenExpiredError),
    (core_exceptions.ForbiddenError, CortexSearchAuthError),
    (core_exceptions.NotFoundError, CortexSearchNotFoundError),
    (core_exceptions.ServerError, CortexSearchServiceError),
    (core_exceptions.ConflictError, CortexSearchError),
    (core_exceptions.RetryTimeoutError, CortexSearchTimeoutError),
    (core_exceptions.InvalidArgumentsError, CortexSearchBadRequestError),
    # network errors
    (ReauthenticationRequest, CortexSearchTokenExpiredError),
]

RETRYABLE_ERRORS = (
    CortexSearchRateLimitError,
    CortexSearchServiceError,
    CortexSearchTimeoutError,
    CortexSearchTokenExpiredError,
)


def map_snowflake_error(exc: Exception) -> CortexSearchError:
    for snowflake_cls, our_cls in _SNOWFLAKE_ERROR_MAP:
        if isinstance(exc, snowflake_cls):
            return our_cls(str(exc))
    return CortexSearchError(str(exc))