from __future__ import annotations

from time import monotonic, sleep
from typing import TYPE_CHECKING, Any

from ..._endpoints import Tasks
from ..._internal._schemas import APIEnvelope
from ..._internal._schemas.images import ArtworkDetails
from ..._internal._schemas.tasks import (
    TaskItem,
    TaskOperationResult,
    TasksCollection,
    TaskSpeedUpResult,
)
from ...exceptions import (
    MissingDataError,
    TaskTimeoutError,
)
from .._base import BaseResource

if TYPE_CHECKING:
    from ..._sync_client import SyncClient
else:
    SyncClient = Any


class TasksResource(BaseResource[SyncClient]):
    """Manages image generation tasks (sync).

    Accessible via ``client.tasks``.
    """

    def batch_progress(
        self,
        task_ids: list[str],
    ) -> APIEnvelope[TasksCollection]:
        """Fetch progress for multiple tasks simultaneously.

        Args:
            task_ids: Task IDs to query.

        Returns:
            ``APIEnvelope[TasksCollection]`` — ``.value.items`` is a list of
            :class:`TaskItem` objects, one per queried task ID.

        Example:
            >>> result = client.tasks.batch_progress(["tid1", "tid2"])
            >>> for item in result.value.items:
            ...     print(item.task_id, item.percent)
        """
        env = self._client.request_route(
            Tasks.PROGRESS,
            json={"task_ids": task_ids},
        )
        return APIEnvelope[TasksCollection].model_validate(env)

    def progress(
        self,
        task_id: str,
    ) -> APIEnvelope[TaskItem]:
        """Fetch progress for a single task.

        Convenience wrapper around :meth:`batch_progress` for a single ID.

        Args:
            task_id: Task ID to query.

        Returns:
            ``APIEnvelope[TaskItem]`` — ``.value.percent`` is 0-100;
            ``.value.is_done`` is ``True`` when the task has finished;
            ``.value.image_url`` / ``.value.image_urls`` hold the result URLs.

        Raises:
            MissingDataError: If the server returns no entry for ``task_id``.

        Example:
            >>> env = client.tasks.progress("tid1")
            >>> print(env.value.percent, env.value.is_done)
        """
        env = self.batch_progress([task_id])
        items = env.value.items
        item = next((i for i in items if i.task_id == task_id), None)
        if item is None:
            raise MissingDataError(
                f"No task progress item returned for task_id={task_id!r}."
            )

        return APIEnvelope[TaskItem](status=env.status, data=item)

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

        Args:
            task_id: Task ID to wait for.
            interval: Seconds to sleep between polling attempts.
            timeout: Maximum seconds to wait before raising. ``None`` waits
                indefinitely.

        Returns:
            The completed :class:`TaskItem`. Use ``.image_url`` /
            ``.image_urls`` to access the generated image(s).

        Raises:
            TaskTimeoutError: If the task does not complete within ``timeout``
                seconds.

        Example:
            >>> item = client.tasks.wait("tid1", timeout=120)
            >>> print(item.image_url)
        """
        deadline = None if timeout is None else (monotonic() + timeout)

        while True:
            env = self.progress(task_id)
            item = env.value

            if item.is_done:
                return item

            if deadline is not None and monotonic() >= deadline:
                raise TaskTimeoutError(
                    f"Task {task_id!r} did not finish within {timeout} seconds."
                )
            sleep(interval)

    def wait_many(
        self,
        task_ids: list[str],
        *,
        interval: float = 2.0,
        timeout: float | None = 300.0,
    ) -> dict[str, TaskItem]:
        """Poll multiple tasks until all finish.

        Args:
            task_ids: Task IDs to wait for. Returns an empty dict immediately
                if the list is empty.
            interval: Seconds to sleep between polling attempts.
            timeout: Maximum seconds to wait before raising. ``None`` waits
                indefinitely.

        Returns:
            A ``dict`` mapping each task ID to its completed :class:`TaskItem`.

        Raises:
            TaskTimeoutError: If any tasks have not completed within ``timeout``
                seconds. The error message lists the remaining task IDs.

        Example:
            >>> results = client.tasks.wait_many(["tid1", "tid2"])
            >>> for task_id, item in results.items():
            ...     print(task_id, item.image_url)
        """
        if not task_ids:
            return {}

        wanted = set(task_ids)
        done: dict[str, TaskItem] = {}

        deadline = None if timeout is None else (monotonic() + timeout)

        while True:
            coll = self.batch_progress(task_ids).value
            for item in coll.items:
                if item.task_id in wanted and item.is_done:
                    done[item.task_id] = item

            if len(done) == len(wanted):
                return done

            if deadline is not None and monotonic() >= deadline:
                remaining = sorted(wanted - done.keys())
                raise TaskTimeoutError(
                    f"Tasks did not finish within {timeout} seconds. Remaining: {remaining}"
                )

            sleep(interval)

    def cancel(
        self,
        task_id: str,
    ) -> APIEnvelope[TaskOperationResult]:
        """Cancel an active generation task.

        Args:
            task_id: Identifier of the task to cancel.

        Returns:
            ``APIEnvelope[TaskOperationResult]`` — ``.value.new_task`` contains
            updated task metadata; ``.value.task_return`` contains queue and
            resource statistics after the cancellation.

        Example:
            >>> result = client.tasks.cancel("tid1")
            >>> print(result.value.new_task.status)
        """
        env = self._client.request_route(
            Tasks.CANCEL,
            json={"task_id": task_id, "action": 5},
        )
        return APIEnvelope[TaskOperationResult].model_validate(env)

    def delete(
        self,
        task_ids: list[str] | str,
    ) -> APIEnvelope[None]:
        """Delete one or more tasks.

        Args:
            task_ids: A single task ID or a list of task IDs to delete.

        Returns:
            ``APIEnvelope[None]`` — the server returns no data payload on
            success.

        Example:
            >>> client.tasks.delete("tid1")
            >>> client.tasks.delete(["tid1", "tid2"])
        """
        return self._client.request_route(
            Tasks.DELETE,
            json={"ids": task_ids if isinstance(task_ids, list) else [task_ids]},
        )

    def speed_up(
        self,
        task_id: str,
    ) -> APIEnvelope[TaskSpeedUpResult]:
        """Request priority processing for a queued task.

        Args:
            task_id: Identifier of the task to accelerate.

        Returns:
            ``APIEnvelope[TaskSpeedUpResult]`` — ``.value.num`` contains the
            remaining number of free speed ups.

        Example:
            >>> result = client.tasks.speed_up("tid1")
            >>> print(result.value.num)
        """
        env = self._client.request_route(
            Tasks.SPEED_UP,
            json={"id": task_id},
        )
        return APIEnvelope[TaskSpeedUpResult].model_validate(env)

    def detail(
        self,
        task_id: str,
        *,
        task_flow_version: str = "v2",
    ) -> APIEnvelope[ArtworkDetails]:
        """Fetch detailed artwork information for a completed task.

        Args:
            task_id: Identifier of the task to look up.
            task_flow_version: Task flow version forwarded to the API.

        Returns:
            ``APIEnvelope[ArtworkDetails]`` — ``.value.img_uris`` contains the
            generated image entries; ``.value.art_model_name`` is the model
            used; ``.value.artwork_no`` is the public artwork identifier.

        Example:
            >>> info = client.tasks.detail("tid1")
            >>> for uri in info.value.img_uris:
            ...     print(uri.url)
        """
        env = self._client.request_route(
            Tasks.DETAIL,
            json={"id": task_id, "task_flow_version": task_flow_version},
        )
        return APIEnvelope[ArtworkDetails].model_validate(env)
