"""Predictions module."""

import asyncio
import time
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Callable, Optional, TypedDict, cast

from wriftai._resource import Resource
from wriftai.common_types import JsonValue, NotRequired, StrEnum, Unpack
from wriftai.pagination import PaginatedResponse, PaginationOptions


class ErrorSource(StrEnum):
    """Enumeration of possible error sources."""

    internal = "internal"
    external = "external"


class Status(StrEnum):
    """Enumeration of possible prediction statuses."""

    pending = "pending"
    started = "started"
    failed = "failed"
    succeeded = "succeeded"
    cancelled = "cancelled"
    abandoned = "abandoned"


class TaskError(TypedDict):
    """Represents an error that occurred during a prediction."""

    source: ErrorSource
    message: str
    detail: JsonValue


class PredictionModel(TypedDict):
    """Represents the prediction model."""

    owner: str
    """Username of the model's owner."""
    name: str
    """Name of the model."""
    version_number: int
    """Version number of the model."""


class Prediction(TypedDict):
    """Represents a prediction."""

    id: str
    """Unique identifier for the prediction."""
    created_at: str
    """Timestamp when the prediction was created."""
    status: Status
    """Current status of the prediction."""
    updated_at: str
    """Timestamp of the last update to the prediction."""
    setup_time: str | None
    """Time taken to set up the prediction environment."""
    execution_time: str | None
    """Time taken to execute the prediction."""
    model: PredictionModel
    """The model of the prediction."""


class PredictionWithIO(Prediction):
    """Represents a prediction with I/O details."""

    input: JsonValue
    """Input data provided for the prediction."""
    output: JsonValue
    """Output generated by the prediction."""
    error: TaskError | None
    """Error details if the prediction failed."""


class Webhook(TypedDict):
    """Represents webhook details."""

    url: str
    """HTTP URL to POST updates to."""
    secret: NotRequired[str | None]
    """Secret to generate the signature for the webhook."""


class CreatePredictionParams(TypedDict):
    """Prediction creation params."""

    model: str
    """model: The model reference in either owner/name
    or owner/name:version-number format (for example:
    deepseek-ai/deepseek-r1 or deepseek-ai/deepseek-r1:1)."""

    input: JsonValue
    """Input data provided for the prediction."""
    webhook: NotRequired[Webhook | None]
    """Details about the webhook to post prediction updates to."""
    validate_input: NotRequired[bool]
    """Enable input validation against the schema before processing."""


@dataclass
class _BaseWaitOptions:
    """Options for customizing wait behaviour."""

    poll_interval: float = 1.0
    """Time in seconds between polling attempts to check the prediction status."""


@dataclass
class WaitOptions(_BaseWaitOptions):
    """Options for customizing wait behaviour."""

    on_poll: Optional[Callable[[PredictionWithIO], None]] = None
    """Optional callback that receives the latest prediction after each poll."""


@dataclass
class AsyncWaitOptions(_BaseWaitOptions):
    """Options for customizing asynchronous wait behaviour."""

    on_poll: Optional[Callable[[PredictionWithIO], Awaitable[None]]] = None
    """Optional callback that receives the latest prediction after each poll."""


DEFAULT_WAIT_OPTIONS = WaitOptions()
"""Default wait options."""

DEFAULT_ASYNC_WAIT_OPTIONS = AsyncWaitOptions()
"""Default async wait options."""


class Predictions(Resource):
    """Resource for operations related to predictions."""

    _API_PREFIX = "/predictions"
    _CANCEL_API_SUFFIX = "/cancel"

    def get(self, prediction_id: str) -> PredictionWithIO:
        """Get a prediction by id.

        Args:
            prediction_id: The unique identifier of the prediction.

        Returns:
            The prediction object.
        """
        return cast(
            PredictionWithIO,
            self._api.request(method="GET", path=f"{self._API_PREFIX}/{prediction_id}"),
        )

    async def async_get(self, prediction_id: str) -> PredictionWithIO:
        """Get a prediction by id.

        Args:
            prediction_id: The unique identifier of the prediction.

        Returns:
            The prediction object.
        """
        return cast(
            PredictionWithIO,
            await self._api.async_request(
                method="GET", path=f"{self._API_PREFIX}/{prediction_id}"
            ),
        )

    def cancel(self, prediction_id: str) -> None:
        """Cancel a prediction by id.

        Args:
            prediction_id: The unique identifier of the prediction.
        """
        self._api.request(
            method="POST",
            path=f"{self._API_PREFIX}/{prediction_id}{self._CANCEL_API_SUFFIX}",
        )

    async def async_cancel(self, prediction_id: str) -> None:
        """Cancel a prediction by id.

        Args:
            prediction_id: The unique identifier of the prediction.
        """
        await self._api.async_request(
            method="POST",
            path=f"{self._API_PREFIX}/{prediction_id}{self._CANCEL_API_SUFFIX}",
        )

    def list(
        self,
        pagination_options: Optional[PaginationOptions] = None,
    ) -> PaginatedResponse[Prediction]:
        """List predictions.

        Args:
            pagination_options: Optional settings to control pagination behavior.

        Returns:
            Paginated response containing predictions and navigation metadata.
        """
        response = self._api.request(
            method="GET", params=pagination_options, path=self._API_PREFIX
        )

        # The response will always match the PaginatedResponse structure,
        # but static type checkers may not infer this correctly.
        # Hence, we ignore the argument type warning.
        return PaginatedResponse(**response)  # type:ignore[arg-type]

    async def async_list(
        self,
        pagination_options: Optional[PaginationOptions] = None,
    ) -> PaginatedResponse[Prediction]:
        """List predictions.

        Args:
            pagination_options: Optional settings to control pagination behavior.

        Returns:
            Paginated response containing predictions and navigation metadata.
        """
        response = await self._api.async_request(
            method="GET", params=pagination_options, path=self._API_PREFIX
        )
        # The response will always match the PaginatedResponse structure,
        # but static type checkers may not infer this correctly.
        # Hence, we ignore the argument type warning.
        return PaginatedResponse(**response)  # type:ignore[arg-type]

    def create(
        self,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
        **params: Unpack[CreatePredictionParams],
    ) -> PredictionWithIO:
        """Create a prediction.

        Args:
            wait: If True, waits until the prediction reaches a terminal state.
                If False, returns the prediction immediately. Defaults to False.
            wait_options: Options for customizing wait behavior. If not provided,
                default options are used.
            params (CreatePredictionParams): Prediction creation params.

        Returns:
            The new prediction.
        """
        headers = None
        if "validate_input" in params:
            headers = {"Validate-Input": str(params.pop("validate_input")).lower()}

        prediction = cast(
            PredictionWithIO,
            self._api.request(
                method="POST",
                path=self._API_PREFIX,
                # The params matches JsonValue at runtime,
                # but static type checkers may not infer this correctly.
                # Hence, we ignore the argument type warning.
                body=params,  # type:ignore[arg-type]
                headers=headers,
            ),
        )

        if wait:
            return self.wait(prediction["id"], options=wait_options)

        return prediction

    async def async_create(
        self,
        wait: bool = False,
        wait_options: AsyncWaitOptions = DEFAULT_ASYNC_WAIT_OPTIONS,
        **params: Unpack[CreatePredictionParams],
    ) -> PredictionWithIO:
        """Create a prediction.

        Args:
            wait: If True, waits until the prediction reaches a terminal state.
                If False, returns the prediction immediately. Defaults to False.
            wait_options: Options for customizing wait behavior. If not provided,
                default options are used.
            params (CreatePredictionParams): Prediction creation params.

        Returns:
            The new prediction.
        """
        headers = None
        if "validate_input" in params:
            headers = {"Validate-Input": str(params.pop("validate_input")).lower()}

        prediction = cast(
            PredictionWithIO,
            await self._api.async_request(
                method="POST",
                path=self._API_PREFIX,
                # The params matches JsonValue at runtime,
                # but static type checkers may not infer this correctly.
                # Hence, we ignore the argument type warning.
                body=params,  # type:ignore[arg-type]
                headers=headers,
            ),
        )

        if wait:
            return await self.async_wait(prediction["id"], options=wait_options)

        return prediction

    @staticmethod
    def _is_status_terminal(status: Status) -> bool:
        """Determines whether the prediction status is terminal or not."""
        return status in {
            Status.succeeded,
            Status.failed,
            Status.cancelled,
            Status.abandoned,
        }

    def wait(
        self,
        predictionID: str,
        options: WaitOptions = DEFAULT_WAIT_OPTIONS,
    ) -> PredictionWithIO:
        """Wait for a prediction to complete.

        This method uses short polling to check the prediction status at
        regular intervals until it reaches a terminal state. It blocks
        execution and waits until the prediction has finished processing.

        Args:
            predictionID: Unique identifier of the prediction.
            options: Options for customizing wait behavior. If not provided,
                default options are used.

        Returns:
            Prediction after completion.
        """
        while True:
            result = self.get(predictionID)
            if options.on_poll:
                options.on_poll(result)
            if Predictions._is_status_terminal(result["status"]):
                return result
            time.sleep(options.poll_interval)

    async def async_wait(
        self,
        predictionID: str,
        options: AsyncWaitOptions = DEFAULT_ASYNC_WAIT_OPTIONS,
    ) -> PredictionWithIO:
        """Wait for a prediction to complete.

        This method uses short polling to check the prediction status at
        regular intervals until it reaches a terminal state. It blocks
        execution and waits until the prediction has finished processing.

        Args:
            predictionID: Unique identifier of the prediction.
            options: Options for customizing wait behavior. If not provided,
                default options are used.

        Returns:
            Prediction after completion.
        """
        while True:
            result = await self.async_get(predictionID)
            if options.on_poll:
                await options.on_poll(result)
            if Predictions._is_status_terminal(result["status"]):
                return result
            await asyncio.sleep(options.poll_interval)
