# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

from __future__ import annotations

import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterator

APPLIER_TRANSIENT_RETRIES = max(
    0, int(os.environ.get("GENEVA_APPLIER_TRANSIENT_RETRIES", "3"))
)
APPLIER_RETRY_BASE_BACKOFF_SECONDS = float(
    os.environ.get("GENEVA_APPLIER_RETRY_BASE_BACKOFF_SECONDS", "1.0")
)
APPLIER_RETRY_MAX_BACKOFF_SECONDS = float(
    os.environ.get("GENEVA_APPLIER_RETRY_MAX_BACKOFF_SECONDS", "8.0")
)

NON_RETRYABLE_OBJECT_STORE_STATUS_CODES = {401, 403, 404}
RETRYABLE_OBJECT_STORE_STATUS_CODES = {408, 429, 500, 502, 503, 504}
OBJECT_STORE_CONTEXT_MARKERS = (
    "aws credentials",
    "credential",
    "object store",
    "lanceerror(io)",
    "connectorerror",
    "providererror",
)
NON_RETRYABLE_OBJECT_STORE_MARKERS = (
    "unauthorized",
    "forbidden",
    "permission denied",
    "access denied",
    "invalid schema",
    "schema mismatch",
    "schema error",
    "missing column",
    "invalid input",
    "invalid argument",
)
RETRYABLE_OBJECT_STORE_MARKERS = (
    "timeout",
    "timed out",
    "temporarily unavailable",
    "temporary failure",
    "temporary redirect",
    "connection reset",
    "connection refused",
    "connection aborted",
    "connection error",
    "broken pipe",
    "unexpected eof",
    "eof while",
    "too many requests",
    "rate limit",
    "throttle",
    "service unavailable",
    "bad gateway",
    "gateway timeout",
)


def get_applier_retry_settings() -> tuple[int, float, float]:
    """Return applier retry settings shared by run/setup object store retries."""
    return (
        APPLIER_TRANSIENT_RETRIES,
        APPLIER_RETRY_BASE_BACKOFF_SECONDS,
        APPLIER_RETRY_MAX_BACKOFF_SECONDS,
    )


def iter_exception_chain(exc: BaseException) -> Iterator[BaseException]:
    """Yield exception and nested causes/contexts without looping."""
    seen: set[int] = set()
    current: BaseException | None = exc
    while current is not None and id(current) not in seen:
        seen.add(id(current))
        yield current
        current = current.__cause__ or current.__context__


def get_status_code(exc: BaseException) -> int | None:
    status_code = getattr(exc, "status_code", None)
    return status_code if isinstance(status_code, int) else None


def is_not_found_object_store_message(msg: str) -> bool:
    return (
        "no such table" in msg
        or "does not exist" in msg
        or ("table" in msg and "not found" in msg)
        or ("column" in msg and "not found" in msg)
        or "not found in schema" in msg
    )


def has_marker(msg: str, markers: tuple[str, ...]) -> bool:
    return any(marker in msg for marker in markers)


def is_retryable_object_store_error(exc: BaseException) -> bool:
    """Return True for transient object-store setup/read failures."""
    candidates = tuple(iter_exception_chain(exc))
    has_object_store_context = False

    for candidate in candidates:
        status_code = get_status_code(candidate)
        if status_code in NON_RETRYABLE_OBJECT_STORE_STATUS_CODES:
            return False

        msg = str(candidate).lower()
        if not msg:
            continue
        if has_marker(msg, OBJECT_STORE_CONTEXT_MARKERS):
            has_object_store_context = True
        if is_not_found_object_store_message(msg):
            return False
        if has_marker(msg, NON_RETRYABLE_OBJECT_STORE_MARKERS):
            return False

    for candidate in candidates:
        status_code = get_status_code(candidate)
        if (
            has_object_store_context
            and status_code in RETRYABLE_OBJECT_STORE_STATUS_CODES
        ):
            return True

        msg = str(candidate).lower()
        if not msg:
            continue
        if has_object_store_context and has_marker(msg, RETRYABLE_OBJECT_STORE_MARKERS):
            return True

    return False
