from __future__ import annotations

from typing import TYPE_CHECKING, Any, Generic, Self, TypeVar, cast

from pydantic import PrivateAttr

from ..._internal._schemas import APIEnvelope

if TYPE_CHECKING:
    from ..._internal._schemas.tasks import TaskItem
    from ..async_.tasks import AsyncTasksResource
    from ..sync.tasks import TasksResource


T = TypeVar("T")


class _BaseTaskEnvelope(APIEnvelope[T], Generic[T]):
    """Base envelope that binds a task response to the tasks resource.

    Extends ``APIEnvelope[T]`` by holding a reference to the tasks resource,
    enabling ``wait()`` and ``progress()`` calls directly on the envelope
    without needing to pass the task ID manually.

    All concrete ``T`` types used with this class must expose an ``id: str``
    attribute.
    """

    _tasks: Any = PrivateAttr()

    @property
    def _task_id(self) -> str:
        val: Any = self.value  # all T used here have `id: str`
        return cast("str", val.id)

    @classmethod
    def from_envelope(cls, envelope: APIEnvelope[T], *, tasks: Any) -> Self:
        """Construct a task envelope from a plain ``APIEnvelope`` and a tasks resource.

        Args:
            envelope: The original API envelope returned by a task-creating endpoint.
            tasks: The tasks resource instance (sync or async) used for polling.

        Returns:
            A new instance of this class with ``_tasks`` set.
        """
        obj = cls.model_validate(envelope.model_dump())
        obj._tasks = tasks
        return obj


class AsyncTaskEnvelope(_BaseTaskEnvelope[T], Generic[T]):
    """Task envelope for use with the async client.

    Returned by task-creating endpoints (e.g. ``client.images.generate``).
    Provides ``await env.wait()`` and ``await env.progress()`` as shortcuts
    so you don't have to manage the task ID yourself.

    Example:
        >>> env = await client.images.generate(...)
        >>> item = await env.wait(timeout=120)
        >>> print(item.image_url)
    """

    _tasks: AsyncTasksResource = PrivateAttr()

    async def wait(
        self, *, interval: float = 2.0, timeout: float | None = 300.0
    ) -> TaskItem:
        """Poll the task until it finishes and return the completed :class:`TaskItem`.

        Args:
            interval: Seconds between polling attempts.
            timeout: Maximum seconds to wait before raising. ``None`` waits
                indefinitely.

        Returns:
            The completed :class:`TaskItem` with image URLs populated.

        Raises:
            TaskTimeoutError: If the task does not complete within ``timeout``
                seconds.
        """
        return await self._tasks.wait(self._task_id, interval=interval, timeout=timeout)

    async def progress(self) -> APIEnvelope[TaskItem]:
        """Fetch the current progress of this task.

        Returns:
            ``APIEnvelope[TaskItem]`` where ``.value.percent`` is 0-100 and
            ``.value.is_done`` indicates completion.

        Raises:
            MissingDataError: If the server returns no entry for this task.
        """
        return await self._tasks.progress(self._task_id)


class SyncTaskEnvelope(_BaseTaskEnvelope[T], Generic[T]):
    """Task envelope for use with the sync client.

    Returned by task-creating endpoints (e.g. ``client.images.generate``).
    Provides ``env.wait()`` and ``env.progress()`` as shortcuts so you don't
    have to manage the task ID yourself.

    Example:
        >>> env = client.images.generate(...)
        >>> item = env.wait(timeout=120)
        >>> print(item.image_url)
    """

    _tasks: TasksResource = PrivateAttr()

    def wait(self, *, interval: float = 2.0, timeout: float | None = 300.0) -> TaskItem:
        """Poll the task until it finishes and return the completed :class:`TaskItem`.

        Args:
            interval: Seconds between polling attempts.
            timeout: Maximum seconds to wait before raising. ``None`` waits
                indefinitely.

        Returns:
            The completed :class:`TaskItem` with image URLs populated.

        Raises:
            TaskTimeoutError: If the task does not complete within ``timeout``
                seconds.
        """
        return self._tasks.wait(self._task_id, interval=interval, timeout=timeout)

    def progress(self) -> APIEnvelope[TaskItem]:
        """Fetch the current progress of this task.

        Returns:
            ``APIEnvelope[TaskItem]`` where ``.value.percent`` is 0-100 and
            ``.value.is_done`` indicates completion.

        Raises:
            MissingDataError: If the server returns no entry for this task.
        """
        return self._tasks.progress(self._task_id)
