from typing import Optional, Set
import grpc
from tenacity import retry, stop_after_attempt, wait_exponential_jitter, retry_if_result
from opentelemetry.trace import Span


class RetryOnRpcErrorClientInterceptor(
    grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor
):
    """
    A gRPC client interceptor that retries RPCs on specific status codes. By default, it retries on UNAVAILABLE and UNKNOWN status codes.

    This interceptor should be placed after the OpenTelemetry interceptor in the interceptor list.
    """

    max_attempts: int
    retryable_status_codes: Set[grpc.StatusCode]

    def __init__(
        self,
        max_attempts: int = 5,
        retryable_status_codes: Set[grpc.StatusCode] = set(
            [grpc.StatusCode.UNAVAILABLE, grpc.StatusCode.UNKNOWN]
        ),
    ) -> None:
        self.max_attempts = max_attempts
        self.retryable_status_codes = retryable_status_codes

    def _intercept_call(self, continuation, client_call_details, request_or_iterator):
        sleep_span: Optional[Span] = None

        def before_sleep(_):
            from chromadb.telemetry.opentelemetry import tracer

            nonlocal sleep_span
            if tracer is not None:
                sleep_span = tracer.start_span("Waiting to retry RPC")

        @retry(
            wait=wait_exponential_jitter(0.1, jitter=0.1),
            stop=stop_after_attempt(self.max_attempts),
            retry=retry_if_result(lambda x: x.code() in self.retryable_status_codes),
            before_sleep=before_sleep,
        )
        def wrapped(*args, **kwargs):
            nonlocal sleep_span
            if sleep_span is not None:
                sleep_span.end()
                sleep_span = None
            return continuation(*args, **kwargs)

        return wrapped(client_call_details, request_or_iterator)

    def intercept_unary_unary(self, continuation, client_call_details, request):
        return self._intercept_call(continuation, client_call_details, request)

    def intercept_unary_stream(self, continuation, client_call_details, request):
        return self._intercept_call(continuation, client_call_details, request)

    def intercept_stream_unary(
        self, continuation, client_call_details, request_iterator
    ):
        return self._intercept_call(continuation, client_call_details, request_iterator)

    def intercept_stream_stream(
        self, continuation, client_call_details, request_iterator
    ):
        return self._intercept_call(continuation, client_call_details, request_iterator)