"""Task result handling and streaming for fair task queue.

This module includes:
- Result handler protocol and implementations
- Task result streaming support
"""

import asyncio
from collections.abc import AsyncIterator
from datetime import UTC, datetime, timedelta
from typing import Any, Protocol

from haolib.background.fair.abstract import FairTask, TaskResultUpdate, TaskStatus

# ============================================================================
# Result Handlers
# ============================================================================


class TaskResultHandler(Protocol):
    """Protocol for handling task results.

    Result handlers are responsible for storing task results after completion.
    This allows the queue to automatically store results from task execution.

    Example:
        ```python
        class StorageResultHandler:
            def __init__(self, storage: AbstractFairQueueStorage):
                self._storage = storage

            async def handle_result(
                self,
                task: FairTask,
                group_id: str,
                result: Any,
            ) -> None:
                await self._storage.store_task_result(
                    task.task_id,
                    group_id,
                    result,
                    ttl_seconds=86400,  # 24 hours
                )
        ```

    """

    async def handle_result(
        self,
        task: FairTask,
        group_id: str,
        result: Any,
    ) -> None:
        """Handle task result after successful completion.

        Args:
            task: Completed task
            group_id: Group ID of the task
            result: Task execution result

        """
        ...


class StorageResultHandler:
    """Result handler that stores results in storage backend.

    This handler automatically stores task results in the storage backend
    so they can be retrieved later via get_task_result().

    Example:
        ```python
        from haolib.background.fair.results import StorageResultHandler

        result_handler = StorageResultHandler(
            storage=storage,
            default_ttl_seconds=86400,  # 24 hours
        )
        queue = FairTaskQueue(..., result_handler=result_handler)
        ```

    """

    def __init__(
        self,
        storage: Any,  # AbstractFairQueueStorage
        default_ttl_seconds: int | None = None,
    ) -> None:
        """Initialize storage result handler.

        Args:
            storage: Storage backend instance
            default_ttl_seconds: Default TTL for stored results (None = use storage default)

        """
        self._storage = storage
        self._default_ttl = default_ttl_seconds

    async def handle_result(
        self,
        task: FairTask,
        group_id: str,
        result: Any,
    ) -> None:
        """Store task result in storage backend.

        Args:
            task: Completed task
            group_id: Group ID of the task
            result: Task execution result

        """
        await self._storage.store_task_result(
            task.task_id,
            group_id,
            result,
            ttl_seconds=self._default_ttl,
        )


class NoOpResultHandler:
    """No-op result handler that does nothing with results.

    Useful for fire-and-forget tasks or when results are handled elsewhere.

    Example:
        ```python
        from haolib.background.fair.results import NoOpResultHandler

        queue = FairTaskQueue(..., result_handler=NoOpResultHandler())
        ```

    """

    async def handle_result(
        self,
        task: FairTask,
        group_id: str,
        result: Any,
    ) -> None:
        """No-op result handler - does nothing.

        Args:
            task: Completed task
            group_id: Group ID of the task
            result: Task execution result

        """


# ============================================================================
# Result Streaming
# ============================================================================


class TaskResultStream:
    """Stream task results as they become available."""

    def __init__(
        self,
        get_status: Any,  # Callable[[str, str | None], Awaitable[TaskStatus | None]]
        get_result: Any,  # Callable[[str, str | None], Awaitable[Any | None]]
        task_id: str,
        group_id: str | None,
        poll_interval: timedelta = timedelta(seconds=0.5),
        timeout_duration: timedelta | None = None,
    ) -> None:
        """Initialize task result stream.

        Args:
            get_status: Async function to get task status
            get_result: Async function to get task result
            task_id: Task ID to stream
            group_id: Optional group ID (will be derived if not provided)
            poll_interval: Interval between status checks
            timeout_duration: Maximum time to wait (None = wait indefinitely)

        """
        self._get_status = get_status
        self._get_result = get_result
        self._task_id = task_id
        self._group_id = group_id
        self._poll_interval = poll_interval
        self._timeout = timeout_duration

    async def stream(self) -> AsyncIterator[TaskResultUpdate]:
        """Stream task status updates and result.

        Yields:
            TaskResultUpdate model with status information

        Raises:
            TimeoutError: If timeout is reached
            RuntimeError: If task fails or is cancelled

        """
        start_time = datetime.now(UTC)

        while True:
            # Check timeout
            if self._timeout:
                elapsed = datetime.now(UTC) - start_time
                if elapsed >= self._timeout:
                    error_msg = f"Task {self._task_id} did not complete within {self._timeout}"
                    raise TimeoutError(error_msg)

            # Get task status
            status = await self._get_status(self._task_id, self._group_id)

            if status is None:
                # Task not found, wait and retry
                await asyncio.sleep(self._poll_interval.total_seconds())
                continue

            # Yield status update
            update = TaskResultUpdate(status=status)

            if status == TaskStatus.COMPLETED:
                # Task completed, get result
                result = await self._get_result(self._task_id, self._group_id)
                if result is not None:
                    update.result = result
                    yield update
                    return
                # Result not yet available, wait a bit
                yield update
                await asyncio.sleep(self._poll_interval.total_seconds())
                continue

            if status == TaskStatus.FAILED:
                update.error = "Task failed"
                yield update
                error_msg = f"Task {self._task_id} failed"
                raise RuntimeError(error_msg)

            if status == TaskStatus.CANCELLED:
                update.error = "Task was cancelled"
                yield update
                error_msg = f"Task {self._task_id} was cancelled"
                raise RuntimeError(error_msg)

            if status == TaskStatus.EXPIRED:
                update.error = "Task expired"
                yield update
                error_msg = f"Task {self._task_id} expired"
                raise RuntimeError(error_msg)

            # Task is still pending or processing, yield status and wait
            yield update
            await asyncio.sleep(self._poll_interval.total_seconds())


async def stream_task_result(
    get_status: Any,  # Callable[[str, str | None], Awaitable[TaskStatus | None]]
    get_result: Any,  # Callable[[str, str | None], Awaitable[Any | None]]
    task_id: str,
    group_id: str | None,
    poll_interval: timedelta = timedelta(seconds=0.5),
    timeout_duration: timedelta | None = None,
) -> AsyncIterator[TaskResultUpdate]:
    """Stream task result updates.

    Args:
        get_status: Async function to get task status
        get_result: Async function to get task result
        task_id: Task ID to stream
        group_id: Optional group ID (will be derived if not provided)
        poll_interval: Interval between status checks
        timeout_duration: Maximum time to wait (None = wait indefinitely)

    Yields:
        TaskResultUpdate model with status information

    Example:
        ```python
        async for update in stream_task_result(
            queue.get_task_status,
            queue.get_task_result,
            task_id,
        ):
            print(f"Status: {update.status}")
            if update.result is not None:
                print(f"Result: {update.result}")
        ```

    """
    stream = TaskResultStream(
        get_status=get_status,
        get_result=get_result,
        task_id=task_id,
        group_id=group_id,
        poll_interval=poll_interval,
        timeout_duration=timeout_duration,
    )
    async for update in stream.stream():
        yield update
