from collections.abc import Callable
from typing import TypeVar, Union

import reactivex
from reactivex import Observable, abc
from reactivex.disposable import SerialDisposable, SingleAssignmentDisposable
from reactivex.internal import curry_flip, is_future
from reactivex.typing import AnyFuture

_T = TypeVar("_T")


def catch_handler(
    source: Observable[_T],
    handler: Callable[
        [Exception, Observable[_T]], Union[Observable[_T], "AnyFuture[_T]"]
    ],
) -> Observable[_T]:
    def subscribe(
        observer: abc.ObserverBase[_T], scheduler: abc.SchedulerBase | None = None
    ) -> abc.DisposableBase:
        d1 = SingleAssignmentDisposable()
        subscription = SerialDisposable()

        subscription.disposable = d1

        def on_error(exception: Exception) -> None:
            try:
                result = handler(exception, source)
            except Exception as ex:  # By design. pylint: disable=W0703
                observer.on_error(ex)
                return

            result = reactivex.from_future(result) if is_future(result) else result
            d = SingleAssignmentDisposable()
            subscription.disposable = d
            d.disposable = result.subscribe(observer, scheduler=scheduler)

        d1.disposable = source.subscribe(
            observer.on_next, on_error, observer.on_completed, scheduler=scheduler
        )
        return subscription

    return Observable(subscribe)


@curry_flip
def catch_(
    source: Observable[_T],
    handler: Observable[_T] | Callable[[Exception, Observable[_T]], Observable[_T]],
) -> Observable[_T]:
    """Continues an observable sequence that is terminated by an
    exception with the next observable sequence.

    Examples:
        >>> res = source.pipe(catch(ys))
        >>> res = catch(ys)(source)
        >>> res = source.pipe(catch(lambda ex, src: ys(ex)))

    Args:
        source: The source observable sequence.
        handler: Second observable sequence used to produce
            results when an error occurred in the first sequence, or an
            exception handler function that returns an observable sequence
            given the error and source observable that occurred in the
            first sequence.

    Returns:
        An observable sequence containing the first sequence's
        elements, followed by the elements of the handler sequence
        in case an exception occurred.
    """
    if callable(handler):
        return catch_handler(source, handler)
    else:
        return reactivex.catch(source, handler)


__all__ = ["catch_"]
