"""Async task waiter for blocking until tasks complete."""

import asyncio
from typing import Optional


class TaskWaiter:
    """
    Allows agents to physically block until a task completes.
    Uses asyncio Futures to suspend execution without CPU usage.

    Usage:
        # Wait for a single task
        completed = await TaskWaiter.wait_for_task("task-123", timeout=300)

        # Wait for any of multiple tasks
        completed_id = await TaskWaiter.wait_for_any(["task-1", "task-2"], timeout=300)

        # Notify waiters when task completes (called by broadcaster)
        await TaskWaiter.notify_completion("task-123")
    """

    # Shared completion futures: task_id -> Future
    _completion_futures: dict[str, asyncio.Future] = {}
    _lock = asyncio.Lock()

    @classmethod
    async def wait_for_task(
        cls,
        task_id: str,
        timeout: Optional[float] = None,
    ) -> bool:
        """
        Block until task completes.

        Args:
            task_id: Task to wait for
            timeout: Max seconds to wait (None = forever)

        Returns:
            True if task completed, False if timeout
        """
        async with cls._lock:
            if task_id not in cls._completion_futures:
                cls._completion_futures[task_id] = asyncio.get_event_loop().create_future()

        future = cls._completion_futures[task_id]

        # If already done, return immediately
        if future.done():
            return True

        try:
            if timeout:
                await asyncio.wait_for(asyncio.shield(future), timeout=timeout)
            else:
                await future
            return True
        except asyncio.TimeoutError:
            return False

    @classmethod
    async def notify_completion(cls, task_id: str) -> None:
        """
        Called when a task completes - wakes up all waiting agents.
        """
        async with cls._lock:
            if task_id in cls._completion_futures:
                future = cls._completion_futures[task_id]
                if not future.done():
                    future.set_result(True)
                # Create new future for any future waiters
                cls._completion_futures[task_id] = asyncio.get_event_loop().create_future()

    @classmethod
    async def wait_for_any(
        cls,
        task_ids: list[str],
        timeout: Optional[float] = None,
    ) -> str | None:
        """
        Wait for ANY of the given tasks to complete.
        Returns the task_id that completed, or None if timeout.
        """
        if not task_ids:
            return None

        async with cls._lock:
            for task_id in task_ids:
                if task_id not in cls._completion_futures:
                    cls._completion_futures[task_id] = asyncio.get_event_loop().create_future()

        futures = {
            asyncio.ensure_future(cls._wait_single(task_id)): task_id
            for task_id in task_ids
        }

        try:
            done, pending = await asyncio.wait(
                futures.keys(),
                timeout=timeout,
                return_when=asyncio.FIRST_COMPLETED,
            )

            # Cancel pending tasks
            for task in pending:
                task.cancel()

            if done:
                # Find which task completed
                for task in done:
                    return futures[task]

            return None
        except asyncio.CancelledError:
            # Cancel all pending
            for task in futures.keys():
                task.cancel()
            raise

    @classmethod
    async def _wait_single(cls, task_id: str) -> str:
        """Helper to wait for a single task and return its ID."""
        await cls.wait_for_task(task_id)
        return task_id

    @classmethod
    async def is_waiting(cls, task_id: str) -> bool:
        """Check if anyone is waiting for a task."""
        async with cls._lock:
            if task_id not in cls._completion_futures:
                return False
            future = cls._completion_futures[task_id]
            # If there's an undone future, someone might be waiting
            return not future.done()

    @classmethod
    async def clear(cls) -> None:
        """Clear all waiting futures (for testing)."""
        async with cls._lock:
            for future in cls._completion_futures.values():
                if not future.done():
                    future.cancel()
            cls._completion_futures.clear()
