"""Synchronization primitives for concurry.

This module provides efficient wait() and gather() functions for managing
multiple futures, following concurrent.futures patterns with enhanced features:
- Adaptive polling algorithms
- Flexible progress tracking
- Nested structure support
- Iterator mode for streaming results
"""

import asyncio
import time
from concurrent.futures import TimeoutError
from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union

from morphic.structs import map_collection

from ..utils.frameworks import _IS_RAY_INSTALLED
from ..utils.progress import ProgressBar
from .constants import PollingAlgorithm, ReturnWhen
from .future import BaseFuture, wrap_future

# Legacy constants for backward compatibility (deprecated)
ALL_COMPLETED = ReturnWhen.ALL_COMPLETED.value
FIRST_COMPLETED = ReturnWhen.FIRST_COMPLETED.value
FIRST_EXCEPTION = ReturnWhen.FIRST_EXCEPTION.value

# Import Ray if available
if _IS_RAY_INSTALLED:
    import ray


# ======================== Helper Functions ========================


def _check_futures_batch(futures_to_check: Iterable[BaseFuture]) -> Set[BaseFuture]:
    """Efficiently check which futures are done.

    For Ray futures: Uses single ray.wait() call for batch checking (efficient IPC)
    For others: Checks .done() individually

    Args:
        futures_to_check: Iterable of BaseFuture instances to check

    Returns:
        Set of completed futures from input
    """
    if len(list(futures_to_check)) == 0:
        return set()

    futures_list = list(futures_to_check)
    completed = set()

    # Optimize Ray futures with batch checking
    if _IS_RAY_INSTALLED:
        # Separate Ray futures from others
        ray_futures = []
        ray_future_map = {}  # Map ObjectRef to BaseFuture wrapper

        for fut in futures_list:
            if hasattr(fut, "_object_ref"):  # RayFuture has _object_ref attribute
                ray_futures.append(fut._object_ref)
                ray_future_map[id(fut._object_ref)] = fut

        # Batch check Ray futures if any exist
        if len(ray_futures) > 0:
            ready, not_ready = ray.wait(ray_futures, num_returns=len(ray_futures), timeout=0)
            for ref in ready:
                if id(ref) in ray_future_map:
                    completed.add(ray_future_map[id(ref)])

    # Check all non-Ray futures individually (and Ray futures not in batch)
    for fut in futures_list:
        if fut not in completed:  # Skip already completed Ray futures
            if fut.done():
                completed.add(fut)

    return completed


def _wrap_all_futures(obj: Any, recurse: bool = False) -> Any:
    """Wrap all future-like objects using wrap_future().

    Args:
        obj: Object or structure containing futures
        recurse: If True, recursively process nested structures

    Returns:
        Structure with all futures wrapped as BaseFuture

    Example:
        >>> futures = [asyncio_future, ray_object_ref, value]
        >>> wrapped = _wrap_all_futures(futures)
        >>> all(isinstance(f, BaseFuture) for f in wrapped)
        True
    """
    if recurse:
        return map_collection(obj, wrap_future, recurse=True)
    else:
        # For non-recursive, only wrap if it's a collection at top level
        if isinstance(obj, (list, tuple, set)):
            collection_type = type(obj)
            return collection_type(wrap_future(item) for item in obj)
        elif isinstance(obj, dict):
            return {k: wrap_future(v) for k, v in obj.items()}
        else:
            return wrap_future(obj)


def _create_progress_tracker(
    progress: Union[bool, Dict, Callable, None], total: Optional[int], desc: str
) -> Optional[Union[ProgressBar, Callable]]:
    """Create progress tracker from progress parameter.

    Args:
        progress: Progress configuration
            - None/False: No progress tracking
            - True: Auto-create ProgressBar with miniters=total/100
            - Dict: Configuration for ProgressBar
            - Callable: Progress callback function
        total: Total number of items
        desc: Description for progress bar

    Returns:
        ProgressBar instance, callable, or None

    Example:
        >>> # Auto progress bar
        >>> pbar = _create_progress_tracker(True, 1000, "Waiting")
        >>> # Custom progress bar
        >>> pbar = _create_progress_tracker({"unit": "task"}, 100, "Processing")
        >>> # Callback function
        >>> callback = _create_progress_tracker(lambda c, t, e: print(f"{c}/{t}"), 10, "")
    """
    if progress is None or progress is False:
        return None

    if callable(progress):
        return progress

    if progress is True:
        # Auto-create with miniters=total/100
        miniters = max(1, total // 100) if total is not None else 1
        return ProgressBar(total=total, desc=desc, miniters=miniters)

    if isinstance(progress, dict):
        # Apply defaults, let user dict override
        config = {"total": total, "desc": desc}
        if total is not None:
            config["miniters"] = max(1, total // 100)
        config.update(progress)  # User settings override defaults
        return ProgressBar(**config)

    raise ValueError(f"Invalid progress type: {type(progress)}. Must be bool, dict, callable, or None")


def _update_progress(
    tracker: Union[ProgressBar, Callable, None],
    completed: int,
    total: int,
    elapsed: float,
) -> None:
    """Update progress tracker.

    Args:
        tracker: ProgressBar, callable, or None
        completed: Number of completed items
        total: Total number of items
        elapsed: Elapsed time in seconds
    """
    if tracker is None:
        return

    if isinstance(tracker, ProgressBar):
        # Update progress bar to current completed count
        current_n = getattr(tracker.pbar, "n", 0)
        delta = completed - current_n
        if delta > 0:
            tracker.update(delta)
    elif callable(tracker):
        # Call progress callback with (completed, total, elapsed)
        tracker(completed, total, elapsed)


# ======================== wait() Function ========================


def wait(
    fs: Union[List, Tuple, Set, Dict, Any],
    *futs,
    timeout: Optional[float] = None,
    return_when: Union[ReturnWhen, str] = ReturnWhen.ALL_COMPLETED,
    polling: Union[PollingAlgorithm, str] = PollingAlgorithm.Adaptive,
    progress: Union[bool, Dict, Callable, None] = None,
    recurse: bool = False,
) -> Tuple[Set[BaseFuture], Set[BaseFuture]]:
    """Wait for futures to complete and return done and not_done sets.

    Follows concurrent.futures.wait() API:
    https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.wait

    This function efficiently waits for futures by:
    - Using adaptive polling to balance responsiveness and CPU usage
    - Batch-checking Ray futures with single ray.wait() call
    - Supporting progress tracking (progress bar or callback)
    - Handling nested structures with recurse parameter

    Args:
        fs: Primary argument - can be:
            - List/tuple/set of futures (most common)
            - Dict with futures as values (keys preserved)
            - Single future
        *futs: Additional futures (only if fs is not a structure)
        timeout: Maximum time to wait in seconds (None = indefinite)
        return_when: Condition for returning. Options:
            - ReturnWhen.ALL_COMPLETED: Wait until all futures are done (default)
            - ReturnWhen.FIRST_COMPLETED: Return when any future completes
            - ReturnWhen.FIRST_EXCEPTION: Return when any future raises exception
            - Can also pass string values: "all_completed", "first_completed", "first_exception"
        polling: Polling algorithm for checking completion
            - PollingAlgorithm enum value or string name
            - Default: Adaptive (adjusts based on completion rate)
        progress: Progress tracking configuration
            - None/False: No progress
            - True: Auto ProgressBar with miniters=len/100
            - Dict: ProgressBar configuration
            - Callable: Progress callback(completed, total, elapsed)
        recurse: If True, recursively wrap nested futures in structures

    Returns:
        Tuple of (done, not_done) sets containing BaseFuture instances

    Raises:
        TimeoutError: If timeout expires before return_when condition met
        ValueError: If return_when is invalid or if fs is a structure and *futs provided

    Example:
        Basic usage (most common):
            ```python
            from concurry import wait, ReturnWhen

            # Wait for list of futures (most common)
            futures = [worker.task(i) for i in range(10)]
            done, not_done = wait(futures, timeout=30)

            # Wait for dict of futures (preserves keys)
            futures_dict = {"task1": f1, "task2": f2}
            done, not_done = wait(futures_dict)

            # Return when first completes
            done, not_done = wait(futures, return_when=ReturnWhen.FIRST_COMPLETED)

            # Or use string
            done, not_done = wait(futures, return_when="first_completed")
            ```

        Multiple individual futures:
            ```python
            # Pass multiple futures directly (less common)
            done, not_done = wait(future1, future2, future3)
            ```

        With progress tracking:
            ```python
            # Auto progress bar
            done, not_done = wait(futures, progress=True)

            # Custom progress bar
            done, not_done = wait(
                futures,
                progress={"desc": "Waiting", "unit": "task", "color": "#ff0000"}
            )

            # Callback function
            def progress_callback(completed, total, elapsed):
                print(f"Progress: {completed}/{total} ({elapsed:.1f}s)")

            done, not_done = wait(futures, progress=progress_callback)
            ```

        With custom polling:
            ```python
            # Use exponential backoff
            done, not_done = wait(futures, polling="exponential")

            # Fixed interval polling
            done, not_done = wait(futures, polling="fixed")
            ```

        With nested structures:
            ```python
            # Nested futures
            nested = [[future1, future2], [future3, future4]]
            done, not_done = wait(nested, recurse=True)
            ```
    """
    # Convert string to ReturnWhen enum if needed
    if isinstance(return_when, str):
        return_when = ReturnWhen(return_when)

    # Validate usage: can't mix structure and *futs
    if len(futs) > 0 and isinstance(fs, (list, tuple, set, dict)):
        raise ValueError(
            "Cannot provide both a structure (list/tuple/set/dict) as first argument "
            "and additional futures via *futs. Either pass a structure, or pass individual futures."
        )

    # Build futures list based on input pattern
    if len(futs) > 0:
        # Multiple individual futures: wait(f1, f2, f3)
        futures_list = [fs] + list(futs)
        futures_list = [wrap_future(f) for f in futures_list]
    elif isinstance(fs, dict):
        # Dict of futures - extract values only, keys are not futures
        futures_list = [wrap_future(v) for v in fs.values()]
    elif isinstance(fs, (list, tuple, set)):
        # Structure of futures
        if recurse:
            wrapped = _wrap_all_futures(fs, recurse=True)
            # Flatten to get all futures
            all_futures = []

            def collect_futures(obj):
                if isinstance(obj, BaseFuture):
                    all_futures.append(obj)
                return obj

            map_collection(wrapped, collect_futures, recurse=True)
            futures_list = all_futures
        else:
            futures_list = [wrap_future(f) for f in fs]
    else:
        # Single future
        futures_list = [wrap_future(fs)]

    # Initialize sets
    done: Set[BaseFuture] = set()
    not_done: Set[BaseFuture] = set(futures_list)

    # Early return if empty
    if len(futures_list) == 0:
        return done, not_done

    # Create polling strategy using factory
    # Import here to avoid circular imports
    from .algorithms.polling import Poller

    if isinstance(polling, str):
        polling = PollingAlgorithm(polling)

    # Poller factory handles config defaults automatically
    strategy = Poller(polling)

    # Create progress tracker
    total = len(futures_list)
    tracker = _create_progress_tracker(progress, total, "Waiting")
    start_time = time.time()

    # Initial check for already-done futures
    initial_done = _check_futures_batch(not_done)
    done.update(initial_done)
    not_done.difference_update(initial_done)
    _update_progress(tracker, len(done), total, time.time() - start_time)

    # Check if we can return early
    if return_when == ReturnWhen.FIRST_COMPLETED and len(done) > 0:
        if isinstance(tracker, ProgressBar):
            tracker.success()
        return done, not_done

    if return_when == ReturnWhen.FIRST_EXCEPTION:
        for fut in done:
            try:
                if fut.exception(timeout=0) is not None:
                    if isinstance(tracker, ProgressBar):
                        tracker.success()
                    return done, not_done
            except Exception:
                pass

    if return_when == ReturnWhen.ALL_COMPLETED and len(not_done) == 0:
        if isinstance(tracker, ProgressBar):
            tracker.success()
        return done, not_done

    # Main polling loop
    while True:
        # Check timeout
        elapsed = time.time() - start_time
        if timeout is not None and elapsed >= timeout:
            if isinstance(tracker, ProgressBar):
                tracker.stop("Timeout")
            raise TimeoutError(
                f"wait() timed out after {elapsed:.2f}s. "
                f"Completed {len(done)}/{total} futures. "
                f"return_when={return_when}"
            )

        # Check batch of not_done futures
        newly_done = _check_futures_batch(not_done)

        if len(newly_done) > 0:
            # Futures completed
            done.update(newly_done)
            not_done.difference_update(newly_done)
            strategy.record_completion()
            _update_progress(tracker, len(done), total, elapsed)

            # Check return conditions
            if return_when == ReturnWhen.FIRST_COMPLETED:
                if isinstance(tracker, ProgressBar):
                    tracker.success()
                return done, not_done

            if return_when == ReturnWhen.FIRST_EXCEPTION:
                for fut in newly_done:
                    try:
                        if fut.exception(timeout=0) is not None:
                            if isinstance(tracker, ProgressBar):
                                tracker.success()
                            return done, not_done
                    except Exception:
                        pass

            if return_when == ReturnWhen.ALL_COMPLETED and len(not_done) == 0:
                if isinstance(tracker, ProgressBar):
                    tracker.success()
                return done, not_done
        else:
            # No completions
            strategy.record_no_completion()

        # Sleep before next check
        interval = strategy.get_next_interval()
        time.sleep(interval)


# ======================== gather() Function ========================


def gather(
    fs: Union[List, Tuple, Set, Dict, Any],
    *futs,
    return_exceptions: bool = False,
    iter: bool = False,
    timeout: Optional[float] = None,
    polling: Union[PollingAlgorithm, str] = PollingAlgorithm.Adaptive,
    progress: Union[bool, Dict, Callable, None] = None,
    recurse: bool = False,
) -> Union[List[Any], Dict[Any, Any], Iterator[Tuple[int, Any]]]:
    """Gather results from multiple futures.

    Similar to asyncio.gather() but works with all future types and provides
    additional features like progress tracking, adaptive polling, and iterator mode.

    Args:
        fs: Primary argument - can be:
            - List/tuple/set of futures (most common) -> returns list
            - Dict with futures as values -> returns dict with same keys
            - Single future -> returns list with one element
        *futs: Additional futures (only if fs is not a structure)
        return_exceptions: If True, return exceptions as results instead of raising
        iter: If True, return generator yielding results as they complete
        timeout: Maximum time to wait for all results (None = indefinite)
        polling: Polling algorithm for checking completion
        progress: Progress tracking configuration
            - None/False: No progress
            - True: Auto ProgressBar with miniters=len/100
            - Dict: ProgressBar configuration
            - Callable: Progress callback(completed, total, elapsed)
        recurse: If True, recursively gather nested futures in structures

    Returns:
        If iter=False and fs is list/tuple: List of results in same order
        If iter=False and fs is dict: Dict with same keys and gathered values
        If iter=True: Generator yielding (index/key, result) as futures complete

    Raises:
        Exception: Any exception from futures (if return_exceptions=False)
        TimeoutError: If timeout expires before all futures complete
        ValueError: If fs is a structure and *futs is provided

    Example:
        Blocking mode (most common):
            ```python
            from concurry import gather

            # Gather list of futures (most common)
            futures = [worker.task(i) for i in range(10)]
            results = gather(futures)  # Returns list

            # Gather dict of futures (preserves keys)
            futures_dict = {"task1": f1, "task2": f2}
            results = gather(futures_dict)  # Returns {"task1": r1, "task2": r2}

            # With exception handling
            results = gather(futures, return_exceptions=True)
            # May contain Exception objects
            ```

        Multiple individual futures (less common):
            ```python
            # Pass multiple futures directly
            results = gather(future1, future2, future3)
            ```

        Iterator mode:
            ```python
            # Process as they arrive (out of order)
            for idx, result in gather(futures, iter=True):
                print(f"Future {idx} completed: {result}")

            # With dict
            futures_dict = {"task1": f1, "task2": f2}
            for key, result in gather(futures_dict, iter=True):
                print(f"{key} completed: {result}")
            ```

        With progress and timeout:
            ```python
            results = gather(
                futures,
                timeout=60,
                progress={"desc": "Gathering", "unit": "result"}
            )
            ```

        Nested structure handling:
            ```python
            # Gather nested futures
            nested_futures = [[f1, f2], [f3, f4]]
            results = gather(nested_futures, recurse=True)
            # Returns: [[r1, r2], [r3, r4]]
            ```
    """
    # Validate usage: can't mix structure and *futs
    if len(futs) > 0 and isinstance(fs, (list, tuple, set, dict)):
        raise ValueError(
            "Cannot provide both a structure (list/tuple/set/dict) as first argument "
            "and additional futures via *futs. Either pass a structure, or pass individual futures."
        )

    # Build args tuple for backend functions
    if len(futs) > 0:
        # Multiple individual futures: gather(f1, f2, f3)
        args = (fs,) + futs
        is_dict_input = False
    else:
        # Single structure or single future
        args = (fs,)
        is_dict_input = isinstance(fs, dict)

    # Delegate to appropriate backend
    if iter:
        return _gather_iter_backend(
            args, return_exceptions, timeout, polling, progress, recurse, is_dict_input
        )
    else:
        return _gather_blocking_backend(
            args, return_exceptions, timeout, polling, progress, recurse, is_dict_input
        )


def _gather_blocking_backend(
    fs: tuple,
    return_exceptions: bool,
    timeout: Optional[float],
    polling: Union[PollingAlgorithm, str],
    progress: Union[bool, Dict, Callable, None],
    recurse: bool,
    is_dict_input: bool,
) -> Union[List[Any], Dict[Any, Any]]:
    """Backend for gather with iter=False. Waits for all futures and returns ordered results.

    Args:
        fs: Tuple of futures (or tuple with single structure)
        return_exceptions: If True, return exceptions as results
        timeout: Maximum wait time
        polling: Polling algorithm
        progress: Progress tracking config
        recurse: Process nested structures
        is_dict_input: If True, first element of fs is a dict, preserve keys

    Returns:
        List of results in same order as input, or dict with keys preserved
    """
    # Handle empty input
    if len(fs) == 0:
        return []

    # Handle dict input specially to preserve keys
    if is_dict_input and len(fs) == 1:
        futures_dict = fs[0]
        if len(futures_dict) == 0:
            return {}

        # Extract keys and futures
        keys = list(futures_dict.keys())
        futures_list = [wrap_future(v) for v in futures_dict.values()]

        # Wait for all futures
        done, not_done = wait(
            futures_list,
            timeout=timeout,
            return_when=ReturnWhen.ALL_COMPLETED,
            polling=polling,
            progress=progress,
        )

        if len(not_done) > 0:
            raise TimeoutError(f"gather() timed out. {len(not_done)}/{len(futures_list)} futures incomplete")

        # Collect results preserving keys
        results_dict = {}
        for key, fut in zip(keys, futures_list):
            try:
                result = fut.result(timeout=0)  # Should be immediate since done
                results_dict[key] = result
            except Exception as e:
                if return_exceptions:
                    results_dict[key] = e
                else:
                    raise

        return results_dict

    # Handle list/tuple/set input (single structure as first arg)
    if len(fs) == 1 and isinstance(fs[0], (list, tuple, set)):
        # Unwrap the structure
        futures_list = [wrap_future(f) for f in fs[0]]
        total = len(futures_list)

        # Wait for all futures
        done, not_done = wait(
            futures_list,
            timeout=timeout,
            return_when=ReturnWhen.ALL_COMPLETED,
            polling=polling,
            progress=progress,
        )

        if len(not_done) > 0:
            raise TimeoutError(f"gather() timed out. {len(not_done)}/{total} futures incomplete")

        # Collect results in order
        results = []
        for fut in futures_list:
            try:
                result = fut.result(timeout=0)  # Should be immediate since done
                results.append(result)
            except Exception as e:
                if return_exceptions:
                    results.append(e)
                else:
                    raise

        return results

    # Wrap all futures
    if recurse:
        wrapped = _wrap_all_futures(fs, recurse=True)
        # For recurse mode, we need to gather recursively through the structure
        # This is complex, so for now we'll gather all futures then reconstruct

        def get_result_or_exception(fut):
            if isinstance(fut, BaseFuture):
                try:
                    return fut.result(timeout=0 if fut.done() else timeout)
                except Exception as e:
                    if return_exceptions:
                        return e
                    else:
                        raise
            return fut

        # Wait for all futures first
        all_futures = []

        def collect_futures(obj):
            if isinstance(obj, BaseFuture):
                all_futures.append(obj)
            return obj

        map_collection(wrapped, collect_futures, recurse=True)

        if len(all_futures) > 0:
            # Wait for all to complete
            done, not_done = wait(
                all_futures,
                timeout=timeout,
                return_when=ReturnWhen.ALL_COMPLETED,
                polling=polling,
                progress=progress,
            )

            if len(not_done) > 0:
                raise TimeoutError(
                    f"gather() timed out. {len(not_done)}/{len(all_futures)} futures incomplete"
                )

        # Now get all results recursively
        results = map_collection(wrapped, get_result_or_exception, recurse=True)
        # If single item in tuple, return as list
        if len(fs) == 1:
            return [results]
        return list(results) if isinstance(results, tuple) else [results]

    else:
        # Non-recursive mode
        futures_list = [wrap_future(f) for f in fs]
        total = len(futures_list)

        # Wait for all futures
        done, not_done = wait(
            futures_list,
            timeout=timeout,
            return_when=ReturnWhen.ALL_COMPLETED,
            polling=polling,
            progress=progress,
        )

        if len(not_done) > 0:
            raise TimeoutError(f"gather() timed out. {len(not_done)}/{total} futures incomplete")

        # Collect results in order
        results = []
        for fut in futures_list:
            try:
                result = fut.result(timeout=0)  # Should be immediate since done
                results.append(result)
            except Exception as e:
                if return_exceptions:
                    results.append(e)
                else:
                    raise

        return results


def _gather_iter_backend(
    fs: tuple,
    return_exceptions: bool,
    timeout: Optional[float],
    polling: Union[PollingAlgorithm, str],
    progress: Union[bool, Dict, Callable, None],
    recurse: bool,
    is_dict_input: bool,
) -> Iterator[Tuple[Union[int, Any], Any]]:
    """Backend generator for gather with iter=True. Yields (index/key, result) as futures complete.

    Args:
        fs: Tuple of futures (or tuple with single structure)
        return_exceptions: If True, yield exceptions as results
        timeout: Maximum wait time
        polling: Polling algorithm
        progress: Progress tracking config
        recurse: Process nested structures
        is_dict_input: If True, first element of fs is a dict, yield keys instead of indices

    Yields:
        (index/key, result) tuples in completion order
    """
    # Handle empty input
    if len(fs) == 0:
        return

    # Handle dict input specially to yield keys
    if is_dict_input and len(fs) == 1:
        futures_dict = fs[0]
        if len(futures_dict) == 0:
            return

        # Extract keys and futures, maintaining mapping
        keys_list = list(futures_dict.keys())
        futures_list = [wrap_future(v) for v in futures_dict.values()]
        future_to_key = {id(fut): key for fut, key in zip(futures_list, keys_list)}
    # Handle list/tuple/set input (single structure as first arg)
    elif len(fs) == 1 and isinstance(fs[0], (list, tuple, set)):
        # Unwrap the structure
        futures_list = [wrap_future(f) for f in fs[0]]
        future_to_key = {id(fut): i for i, fut in enumerate(futures_list)}
    else:
        # Wrap all futures
        if recurse:
            wrapped = _wrap_all_futures(fs, recurse=True)
            all_futures = []
            future_to_path = {}  # Map future to its path in structure

            def collect_with_path(obj, path=[]):
                if isinstance(obj, BaseFuture):
                    all_futures.append(obj)
                    future_to_path[id(obj)] = path.copy()
                return obj

            # Collect all futures with their paths
            def recursive_collect(obj, path=[]):
                if isinstance(obj, BaseFuture):
                    collect_with_path(obj, path)
                elif isinstance(obj, (list, tuple)):
                    for i, item in enumerate(obj):
                        recursive_collect(item, path + [i])
                elif isinstance(obj, dict):
                    for k, v in obj.items():
                        recursive_collect(v, path + [k])
                return obj

            recursive_collect(wrapped)

            futures_list = all_futures
            future_to_key = future_to_path
        else:
            futures_list = [wrap_future(f) for f in fs]
            future_to_key = {id(fut): i for i, fut in enumerate(futures_list)}

    total = len(futures_list)
    if total == 0:
        return

    # Create polling strategy using factory
    # Import here to avoid circular imports
    from .algorithms.polling import Poller

    if isinstance(polling, str):
        polling = PollingAlgorithm(polling)

    # Poller factory handles config defaults automatically
    strategy = Poller(polling)

    # Create progress tracker
    tracker = _create_progress_tracker(progress, total, "Gathering")
    start_time = time.time()

    # Track yielded futures
    pending = set(futures_list)
    yielded_count = 0

    # Main loop
    while len(pending) > 0:
        # Check timeout
        elapsed = time.time() - start_time
        if timeout is not None and elapsed >= timeout:
            if isinstance(tracker, ProgressBar):
                tracker.stop("Timeout")
            raise TimeoutError(
                f"gather(iter=True) timed out after {elapsed:.2f}s. Yielded {yielded_count}/{total} results"
            )

        # Check batch
        newly_done = _check_futures_batch(pending)

        if len(newly_done) > 0:
            # Yield completed futures
            for fut in newly_done:
                # Get the index or key for this future
                key_or_index = future_to_key[id(fut)]
                # For non-recurse non-dict, it's just an index
                # For dict, it's the actual key
                # For recurse, it could be a path (list)
                if isinstance(key_or_index, list) and len(key_or_index) == 1:
                    key_or_index = key_or_index[0]

                try:
                    result = fut.result(timeout=0)
                    yield (key_or_index, result)
                except Exception as e:
                    if return_exceptions:
                        yield (key_or_index, e)
                    else:
                        if isinstance(tracker, ProgressBar):
                            tracker.failure()
                        raise

                yielded_count += 1

            pending.difference_update(newly_done)
            strategy.record_completion()
            _update_progress(tracker, yielded_count, total, elapsed)
        else:
            strategy.record_no_completion()

        # Sleep before next check
        if len(pending) > 0:
            interval = strategy.get_next_interval()
            time.sleep(interval)

    # All done
    if isinstance(tracker, ProgressBar):
        tracker.success()


# ======================== async_wait() Function ========================


async def async_wait(
    fs: Union[List, Tuple, Set, Dict, Any],
    *futs,
    timeout: Optional[float] = None,
    return_when: Union[ReturnWhen, str] = ReturnWhen.ALL_COMPLETED,
    progress: Union[bool, Dict, Callable, None] = None,
) -> Tuple[Set[Any], Set[Any]]:
    """Async version of wait() for use in async contexts.

    This function is designed to be used with coroutines and asyncio tasks. It uses
    native asyncio.wait() for efficient waiting without polling, and properly yields
    control to the event loop.

    Args:
        fs: Primary argument - can be:
            - List/tuple/set of coroutines, asyncio tasks, or futures
            - Dict with coroutines/tasks/futures as values
            - Single coroutine, task, or future
        *futs: Additional coroutines/tasks/futures (only if fs is not a structure)
        timeout: Maximum time to wait in seconds (None = indefinite)
        return_when: Condition for returning. Options:
            - ReturnWhen.ALL_COMPLETED: Wait until all complete (default)
            - ReturnWhen.FIRST_COMPLETED: Return when any completes
            - ReturnWhen.FIRST_EXCEPTION: Return when any raises exception
            - Can also pass string values: "all_completed", "first_completed", "first_exception"
        progress: Progress tracking configuration
            - None/False: No progress
            - True: Auto ProgressBar with miniters=len/100
            - Dict: ProgressBar configuration
            - Callable: Progress callback(completed, total, elapsed)

    Returns:
        Tuple of (done, not_done) sets containing the original coroutines/tasks/futures

    Raises:
        TimeoutError: If timeout expires before return_when condition met
        ValueError: If return_when is invalid or if fs is a structure and *futs provided

    Example:
        Basic usage:
            ```python
            import asyncio
            from concurry import async_wait, ReturnWhen

            async def process():
                async def fetch(url):
                    await asyncio.sleep(0.1)
                    return f"data from {url}"

                # Create coroutines
                coros = [fetch(f"https://api.example.com/{i}") for i in range(10)]

                # Wait for all to complete
                done, not_done = await async_wait(coros, timeout=30)

                # Get results
                results = [await f for f in done]
                return results

            # Run the async function
            asyncio.run(process())
            ```

        With progress tracking:
            ```python
            async def process_with_progress():
                coros = [fetch(url) for url in urls]
                done, not_done = await async_wait(
                    coros,
                    progress={"desc": "Fetching", "unit": "request"}
                )
            ```

        Return on first completion:
            ```python
            async def race_example():
                coros = [task1(), task2(), task3()]
                done, not_done = await async_wait(
                    coros,
                    return_when=ReturnWhen.FIRST_COMPLETED
                )
                # Cancel remaining tasks
                for task in not_done:
                    task.cancel()
            ```

        With Worker futures (use regular wait() instead):
            ```python
            # ✅ Correct: Use regular wait() with Worker futures
            from concurry import Worker, wait

            class AsyncWorker(Worker):
                async def fetch(self, url):
                    await asyncio.sleep(0.1)
                    return f"data from {url}"

            worker = AsyncWorker.options(mode="asyncio").init()
            futures = [worker.fetch(f"url_{i}") for i in range(10)]

            # Use regular wait(), not async_wait()
            done, not_done = wait(futures, timeout=30)
            results = [f.result() for f in done]
            worker.stop()

            # ❌ Wrong: Don't use async_wait() with Worker futures
            # await async_wait(futures)  # Won't work correctly
            ```

    Note:
        - This function uses a polling loop to enable dynamic progress updates.
        - The poll interval defaults to global_config.defaults.async_wait_poll_interval
          (100 microseconds = 10,000 checks/sec).
        - Use async_wait() for raw coroutines in async contexts.
        - Use regular wait() for Worker futures (works across all execution modes).
    """
    # Get poll interval from config
    from ..config import global_config

    local_config = global_config.clone()
    poll_interval = local_config.defaults.async_wait_poll_interval

    # Convert string to ReturnWhen enum if needed
    if isinstance(return_when, str):
        return_when = ReturnWhen(return_when)

    # Validate usage: can't mix structure and *futs
    if len(futs) > 0 and isinstance(fs, (list, tuple, set, dict)):
        raise ValueError(
            "Cannot provide both a structure (list/tuple/set/dict) as first argument "
            "and additional futures via *futs. Either pass a structure, or pass individual futures."
        )

    # Build list of awaitables based on input pattern
    if len(futs) > 0:
        # Multiple individual items: async_wait(c1, c2, c3)
        awaitables_list = [fs] + list(futs)
    elif isinstance(fs, dict):
        # Dict - extract values only, keys are not awaitables
        awaitables_list = list(fs.values())
    elif isinstance(fs, (list, tuple, set)):
        # Structure of awaitables
        awaitables_list = list(fs)
    else:
        # Single awaitable
        awaitables_list = [fs]

    # Ensure all are tasks (convert awaitables to tasks, wrap non-awaitables)
    tasks = []
    for item in awaitables_list:
        if asyncio.iscoroutine(item) or asyncio.isfuture(item):
            tasks.append(asyncio.ensure_future(item))
        else:
            # Non-awaitable value - wrap in completed task
            future = asyncio.get_running_loop().create_future()
            future.set_result(item)
            tasks.append(future)
    total = len(tasks)

    # Early return if empty
    if total == 0:
        return set(), set()

    # Create progress tracker
    tracker = _create_progress_tracker(progress, total, "Waiting")
    start_time = time.time()

    # Map return_when to asyncio constants
    asyncio_return_when = {
        ReturnWhen.ALL_COMPLETED: asyncio.ALL_COMPLETED,
        ReturnWhen.FIRST_COMPLETED: asyncio.FIRST_COMPLETED,
        ReturnWhen.FIRST_EXCEPTION: asyncio.FIRST_EXCEPTION,
    }[return_when]

    # Initial check for already-done tasks
    done_set = {t for t in tasks if t.done()}
    not_done_set = set(tasks) - done_set
    _update_progress(tracker, len(done_set), total, time.time() - start_time)

    # Check if we can return early
    if return_when == ReturnWhen.FIRST_COMPLETED and len(done_set) > 0:
        if isinstance(tracker, ProgressBar):
            tracker.success()
        return done_set, not_done_set

    if return_when == ReturnWhen.FIRST_EXCEPTION:
        for task in done_set:
            try:
                if task.exception() is not None:
                    if isinstance(tracker, ProgressBar):
                        tracker.success()
                    return done_set, not_done_set
            except (asyncio.CancelledError, asyncio.InvalidStateError):
                pass

    if return_when == ReturnWhen.ALL_COMPLETED and len(not_done_set) == 0:
        if isinstance(tracker, ProgressBar):
            tracker.success()
        return done_set, not_done_set

    # Polling loop with progress updates
    try:
        while True:
            elapsed = time.time() - start_time

            # Check timeout
            if timeout is not None and elapsed >= timeout:
                if isinstance(tracker, ProgressBar):
                    tracker.stop("Timeout")
                raise TimeoutError(
                    f"async_wait() timed out after {elapsed:.2f}s. "
                    f"Completed {len(done_set)}/{total} tasks. "
                    f"return_when={return_when}"
                )

            # Check which tasks are done
            newly_done = {t for t in not_done_set if t.done()}
            if len(newly_done) > 0:
                done_set.update(newly_done)
                not_done_set.difference_update(newly_done)
                _update_progress(tracker, len(done_set), total, elapsed)

                # Check return conditions
                if return_when == ReturnWhen.FIRST_COMPLETED:
                    if isinstance(tracker, ProgressBar):
                        tracker.success()
                    return done_set, not_done_set

                if return_when == ReturnWhen.FIRST_EXCEPTION:
                    for task in newly_done:
                        try:
                            if task.exception() is not None:
                                if isinstance(tracker, ProgressBar):
                                    tracker.success()
                                return done_set, not_done_set
                        except (asyncio.CancelledError, asyncio.InvalidStateError):
                            pass

                if return_when == ReturnWhen.ALL_COMPLETED and len(not_done_set) == 0:
                    if isinstance(tracker, ProgressBar):
                        tracker.success()
                    return done_set, not_done_set

            # Sleep before next check (yield to event loop)
            remaining_timeout = None if timeout is None else timeout - elapsed
            sleep_time = (
                min(poll_interval, remaining_timeout) if remaining_timeout is not None else poll_interval
            )
            if sleep_time > 0:
                await asyncio.sleep(sleep_time)

    except asyncio.TimeoutError:
        elapsed = time.time() - start_time
        if isinstance(tracker, ProgressBar):
            tracker.stop("Timeout")

        raise TimeoutError(
            f"async_wait() timed out after {elapsed:.2f}s. "
            f"Completed {len(done_set)}/{total} tasks. "
            f"return_when={return_when}"
        )


# ======================== async_gather() Function ========================


async def async_gather(
    fs: Union[List, Tuple, Set, Dict, Any],
    *futs,
    return_exceptions: bool = False,
    timeout: Optional[float] = None,
    progress: Union[bool, Dict, Callable, None] = None,
) -> Union[List[Any], Dict[Any, Any]]:
    """Async version of gather() for use in async contexts.

    This function is designed to be used with coroutines and asyncio tasks. It uses
    native asyncio.gather() for efficient gathering without polling, and properly
    yields control to the event loop.

    Args:
        fs: Primary argument - can be:
            - List/tuple/set of coroutines, asyncio tasks, or futures
            - Dict with coroutines/tasks/futures as values -> returns dict with same keys
            - Single coroutine, task, or future
        *futs: Additional coroutines/tasks/futures (only if fs is not a structure)
        return_exceptions: If True, return exceptions as results instead of raising
        timeout: Maximum time to wait for all results (None = indefinite)
        progress: Progress tracking configuration
            - None/False: No progress
            - True: Auto ProgressBar with miniters=len/100
            - Dict: ProgressBar configuration
            - Callable: Progress callback(completed, total, elapsed)

    Returns:
        If fs is list/tuple/set: List of results in same order as input
        If fs is dict: Dict with same keys and gathered values

    Raises:
        Exception: Any exception from tasks (if return_exceptions=False)
        TimeoutError: If timeout expires before all tasks complete
        ValueError: If fs is a structure and *futs is provided

    Example:
        Basic usage:
            ```python
            import asyncio
            from concurry import async_gather

            async def batch_process():
                async def compute(x):
                    await asyncio.sleep(0.01)
                    return x ** 2

                # Create coroutines
                coros = [compute(i) for i in range(100)]

                # Gather all results
                results = await async_gather(coros, timeout=60)
                return results

            # Run the async function
            results = asyncio.run(batch_process())
            ```

        With dict (preserves keys):
            ```python
            async def fetch_multiple():
                async def fetch(name):
                    await asyncio.sleep(0.1)
                    return f"data_{name}"

                coros = {
                    "user": fetch("user"),
                    "posts": fetch("posts"),
                    "comments": fetch("comments")
                }

                # Returns dict with same keys
                results = await async_gather(coros)
                # {'user': 'data_user', 'posts': 'data_posts', 'comments': 'data_comments'}
                return results
            ```

        With exception handling:
            ```python
            async def handle_errors():
                async def may_fail(x):
                    if x == 5:
                        raise ValueError("Failed at 5")
                    return x * 2

                coros = [may_fail(i) for i in range(10)]

                # Return exceptions as values
                results = await async_gather(coros, return_exceptions=True)

                # Process results and errors
                for i, result in enumerate(results):
                    if isinstance(result, Exception):
                        print(f"Task {i} failed: {result}")
                    else:
                        print(f"Task {i} succeeded: {result}")
            ```

        With progress tracking:
            ```python
            async def process_with_progress():
                coros = [fetch(url) for url in urls]
                results = await async_gather(
                    coros,
                    timeout=120,
                    progress={"desc": "Fetching", "unit": "request", "color": "#00ff00"}
                )
                return results
            ```

        Mixed types:
            ```python
            async def mixed_example():
                from concurrent.futures import ThreadPoolExecutor

                async def async_task(x):
                    await asyncio.sleep(0.01)
                    return x * 2

                # Mix coroutines, regular values, and thread futures
                with ThreadPoolExecutor() as executor:
                    items = [
                        async_task(10),  # Coroutine
                        42,              # Regular value
                        executor.submit(lambda: 99)  # Thread future
                    ]

                    results = await async_gather(items)
                    # [20, 42, 99]
            ```

        With Worker futures (use regular gather() instead):
            ```python
            # ✅ Correct: Use regular gather() with Worker futures
            from concurry import Worker, gather
            import asyncio

            class AsyncWorker(Worker):
                async def compute(self, x):
                    await asyncio.sleep(0.01)
                    return x ** 2

            worker = AsyncWorker.options(mode="asyncio").init()
            futures = [worker.compute(i) for i in range(10)]

            # Use regular gather(), not async_gather()
            results = gather(futures, timeout=30.0)
            # [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

            worker.stop()

            # ❌ Wrong: Don't use async_gather() with Worker futures
            # results = await async_gather(futures)  # Won't work correctly
            ```

    Note:
        - This function uses a polling loop to enable dynamic progress updates.
        - The poll interval defaults to global_config.defaults.async_gather_poll_interval
          (100 microseconds = 10,000 checks/sec).
        - Use async_gather() for raw coroutines in async contexts.
        - Use regular gather() for Worker futures (works across all execution modes).
    """
    # Get poll interval from config
    from ..config import global_config

    local_config = global_config.clone()
    poll_interval = local_config.defaults.async_gather_poll_interval

    # Validate usage: can't mix structure and *futs
    if len(futs) > 0 and isinstance(fs, (list, tuple, set, dict)):
        raise ValueError(
            "Cannot provide both a structure (list/tuple/set/dict) as first argument "
            "and additional futures via *futs. Either pass a structure, or pass individual futures."
        )

    # Determine if dict input (to preserve keys)
    is_dict_input = False
    keys = None

    # Build list of awaitables based on input pattern
    if len(futs) > 0:
        # Multiple individual items: async_gather(c1, c2, c3)
        awaitables_list = [fs] + list(futs)
    elif isinstance(fs, dict):
        # Dict - extract keys and values
        is_dict_input = True
        keys = list(fs.keys())
        awaitables_list = list(fs.values())
    elif isinstance(fs, (list, tuple, set)):
        # Structure of awaitables
        awaitables_list = list(fs)
    else:
        # Single awaitable
        awaitables_list = [fs]

    total = len(awaitables_list)

    # Early return if empty
    if total == 0:
        return {} if is_dict_input else []

    # Create progress tracker
    tracker = _create_progress_tracker(progress, total, "Gathering")
    start_time = time.time()

    # Ensure all are tasks (convert awaitables to tasks, wrap non-awaitables)
    tasks = []
    for item in awaitables_list:
        if asyncio.iscoroutine(item) or asyncio.isfuture(item):
            tasks.append(asyncio.ensure_future(item))
        else:
            # Non-awaitable value - wrap in completed task
            future = asyncio.get_running_loop().create_future()
            future.set_result(item)
            tasks.append(future)

    # Initial check for already-done tasks
    completed_count = sum(1 for t in tasks if t.done())
    _update_progress(tracker, completed_count, total, time.time() - start_time)

    # If all tasks are already done, gather results immediately
    if completed_count == total:
        if isinstance(tracker, ProgressBar):
            tracker.success()
        results = []
        for task in tasks:
            try:
                result = task.result()
                results.append(result)
            except Exception as e:
                if return_exceptions:
                    results.append(e)
                else:
                    raise
        if is_dict_input:
            return {k: v for k, v in zip(keys, results)}
        else:
            return results

    # Polling loop with progress updates
    try:
        while True:
            elapsed = time.time() - start_time

            # Check timeout
            if timeout is not None and elapsed >= timeout:
                if isinstance(tracker, ProgressBar):
                    tracker.stop("Timeout")
                # Cancel remaining tasks
                for task in tasks:
                    if not task.done():
                        task.cancel()

                completed = sum(1 for t in tasks if t.done())
                raise TimeoutError(
                    f"async_gather() timed out after {elapsed:.2f}s. Completed {completed}/{total} tasks."
                )

            # Check task completion
            newly_completed = sum(1 for t in tasks if t.done())
            if newly_completed > completed_count:
                completed_count = newly_completed
                _update_progress(tracker, completed_count, total, elapsed)

            # All tasks completed
            if completed_count == total:
                if isinstance(tracker, ProgressBar):
                    tracker.success()

                # Gather results from completed tasks
                results = []
                for task in tasks:
                    try:
                        result = task.result()
                        results.append(result)
                    except Exception as e:
                        if return_exceptions:
                            results.append(e)
                        else:
                            if isinstance(tracker, ProgressBar):
                                tracker.failure()
                            raise

                # Return dict if input was dict
                if is_dict_input:
                    return {k: v for k, v in zip(keys, results)}
                else:
                    return results

            # Sleep before next check (yield to event loop)
            remaining_timeout = None if timeout is None else timeout - elapsed
            sleep_time = (
                min(poll_interval, remaining_timeout) if remaining_timeout is not None else poll_interval
            )
            if sleep_time > 0:
                await asyncio.sleep(sleep_time)

    except asyncio.TimeoutError:
        elapsed = time.time() - start_time
        if isinstance(tracker, ProgressBar):
            tracker.stop("Timeout")

        # Cancel remaining tasks
        for task in tasks:
            if not task.done():
                task.cancel()

        # Count completed
        completed = sum(1 for t in tasks if t.done())

        raise TimeoutError(
            f"async_gather() timed out after {elapsed:.2f}s. Completed {completed}/{total} tasks."
        )
