from collections import OrderedDict
from collections.abc import Callable
from typing import Any, TypeVar

from reactivex import Observable, abc
from reactivex.disposable import CompositeDisposable, SingleAssignmentDisposable
from reactivex.internal import noop
from reactivex.operators import take

_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")


def join_(
    right: Observable[_T2],
    left_duration_mapper: Callable[[Any], Observable[Any]],
    right_duration_mapper: Callable[[Any], Observable[Any]],
) -> Callable[[Observable[_T1]], Observable[tuple[_T1, _T2]]]:
    def join(source: Observable[_T1]) -> Observable[tuple[_T1, _T2]]:
        """Correlates the elements of two sequences based on
        overlapping durations.

        Args:
            source: Source observable.

        Return:
            An observable sequence that contains elements
            combined into a tuple from source elements that have an overlapping
            duration.
        """

        left = source

        def subscribe(
            observer: abc.ObserverBase[tuple[_T1, _T2]],
            scheduler: abc.SchedulerBase | None = None,
        ) -> abc.DisposableBase:
            group = CompositeDisposable()
            left_done = False
            left_map: OrderedDict[int, _T1] = OrderedDict()
            left_id = 0
            right_done = False
            right_map: OrderedDict[int, _T2] = OrderedDict()
            right_id = 0

            def on_next_left(value: _T1):
                nonlocal left_id
                duration = None
                current_id = left_id
                left_id += 1
                md = SingleAssignmentDisposable()

                left_map[current_id] = value
                group.add(md)

                def expire():
                    if current_id in left_map:
                        del left_map[current_id]
                    if not len(left_map) and left_done:
                        observer.on_completed()

                    group.remove(md)

                try:
                    duration = left_duration_mapper(value)
                except Exception as exception:
                    observer.on_error(exception)
                    return

                md.disposable = duration.pipe(take(1)).subscribe(
                    noop, observer.on_error, lambda: expire(), scheduler=scheduler
                )

                for val in right_map.values():
                    result = (value, val)
                    observer.on_next(result)

            def on_completed_left() -> None:
                nonlocal left_done
                left_done = True
                if right_done or not len(left_map):
                    observer.on_completed()

            group.add(
                left.subscribe(
                    on_next_left,
                    observer.on_error,
                    on_completed_left,
                    scheduler=scheduler,
                )
            )

            def on_next_right(value: _T2):
                nonlocal right_id
                duration = None
                current_id = right_id
                right_id += 1
                md = SingleAssignmentDisposable()
                right_map[current_id] = value
                group.add(md)

                def expire():
                    if current_id in right_map:
                        del right_map[current_id]
                    if not len(right_map) and right_done:
                        observer.on_completed()

                    group.remove(md)

                try:
                    duration = right_duration_mapper(value)
                except Exception as exception:
                    observer.on_error(exception)
                    return

                md.disposable = duration.pipe(take(1)).subscribe(
                    noop, observer.on_error, lambda: expire(), scheduler=scheduler
                )

                for val in left_map.values():
                    result = (val, value)
                    observer.on_next(result)

            def on_completed_right():
                nonlocal right_done
                right_done = True
                if left_done or not len(right_map):
                    observer.on_completed()

            group.add(
                right.subscribe(
                    on_next_right,
                    observer.on_error,
                    on_completed_right,
                    scheduler=scheduler,
                )
            )
            return group

        return Observable(subscribe)

    return join


__all__ = ["join_"]
