"""Predictions module."""

import asyncio
import time
from dataclasses import dataclass
from typing import Optional, TypedDict, cast, overload

from wriftai._resource import Resource
from wriftai.common_types import JsonValue, NotRequired, StrEnum
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"


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

    source: ErrorSource
    message: str
    detail: JsonValue


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

    url: str
    """URL to access the prediction resource."""
    id: str
    """Unique identifier for the prediction."""
    version_id: str
    """Identifier of the model version used."""
    created_at: str
    """Timestamp when the prediction was created."""
    status: Status
    """Current status of the prediction."""
    webhook_url: str | None
    """URL for webhook callbacks, if any."""
    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."""
    hardware_id: str
    """Identifier of the hardware used for prediction."""
    error: TaskError | None
    """Error details if the prediction failed."""


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."""
    logs: list[str] | None
    """Logs generated during prediction execution."""
    setup_logs: list[str] | None
    """Logs generated during setup."""


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

    input: JsonValue
    """Input data provided for the prediction."""
    webhook_url: NotRequired[str]
    """An HTTP URL to POST prediction updates to."""
    validate_input: NotRequired[bool]
    """Enable input validation against the schema before processing."""


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

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


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


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

    _API_PREFIX = "/predictions"
    _PREDICTIONS_API_SUFFIX = "/predictions"
    _ERROR_MSG_INVALID_PREDICTION_PARAMS = (
        "Either a version_id or both model_owner and model_name must be provided."
    )

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

        Args:
            prediction_id (str): The unique identifier of the prediction.

        Returns:
            PredictionWithIO: 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 (str): The unique identifier of the prediction.

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

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

        Args:
            pagination_options (Optional[PaginationOptions]): Optional settings
                to control pagination behavior.

        Returns:
            PaginatedResponse[Prediction]: 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[PaginationOptions]): Optional settings
                to control pagination behavior.

        Returns:
            PaginatedResponse[Prediction]: 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 _prediction_path(
        self,
        model_owner: Optional[str] = None,
        model_name: Optional[str] = None,
        version_id: Optional[str] = None,
    ) -> str:
        """Constructs the API path to create a prediction.

        Args:
            model_owner (Optional[str]): Username of the model's owner.
            model_name (Optional[str]): Name of the model.
            version_id (Optional[str]): Unique identifier of the model's version.

        Returns:
            str: The API path.

        Raises:
            TypeError: When an invalid combination of version_id, model_owner,
                and model_name is provided.
        """
        if version_id is None and model_owner is not None and model_name is not None:
            return (
                f"{self._MODELS_API_PREFIX}/{model_owner}/"
                f"{model_name}{self._PREDICTIONS_API_SUFFIX}"
            )
        elif model_owner is None and model_name is None and version_id is not None:
            return (
                f"{self._VERSIONS_API_PREFIX}/{version_id}"
                f"{self._PREDICTIONS_API_SUFFIX}"
            )
        else:
            raise TypeError(self._ERROR_MSG_INVALID_PREDICTION_PARAMS)

    @overload
    def create(
        self,
        *,
        model_owner: str,
        model_name: str,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
        params: CreatePredictionParams,
    ) -> PredictionWithIO: ...

    @overload
    def create(
        self,
        *,
        version_id: str,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
        params: CreatePredictionParams,
    ) -> PredictionWithIO: ...

    def create(
        self,
        params: CreatePredictionParams,
        model_owner: Optional[str] = None,
        model_name: Optional[str] = None,
        version_id: Optional[str] = None,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
    ) -> PredictionWithIO:
        """Create a prediction.

        This method creates a prediction using either the latest version of a model
        or a specific model version.

        Args:
            model_owner (`Optional[str]`): Username of the model's owner.
            model_name (`Optional[str]`): Name of the model.
            version_id (`Optional[str]`): Unique identifier of the model's version.
            params (CreatePredictionParams): Prediction creation params.
            wait (bool): If True, waits until the prediction reaches a terminal state.
                If False, returns the prediction immediately. Defaults to False.
            wait_options (WaitOptions): Options for customizing wait behavior.
                If not provided, default options are used.

        Returns:
            PredictionWithIO: The new prediction.

        Raises:
            TypeError: When an invalid combination of version_id, model_owner,
                and model_name is provided.
        """
        path = self._prediction_path(model_owner, model_name, version_id)

        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=path,
                # 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

    @overload
    async def async_create(
        self,
        *,
        model_owner: str,
        model_name: str,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
        params: CreatePredictionParams,
    ) -> PredictionWithIO: ...

    @overload
    async def async_create(
        self,
        *,
        version_id: str,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
        params: CreatePredictionParams,
    ) -> PredictionWithIO: ...

    async def async_create(
        self,
        params: CreatePredictionParams,
        model_owner: Optional[str] = None,
        model_name: Optional[str] = None,
        version_id: Optional[str] = None,
        wait: bool = False,
        wait_options: WaitOptions = DEFAULT_WAIT_OPTIONS,
    ) -> PredictionWithIO:
        """Create a prediction.

        This method creates a prediction using either the latest version of a model
        or a specific model version.

        Args:
            model_owner (`Optional[str]`): Username of the model's owner.
            model_name (`Optional[str]`): Name of the model.
            version_id (`Optional[str]`): Unique identifier of the model's version.
            params (CreatePredictionParams): Prediction creation params.
            wait (bool): If True, waits until the prediction reaches a terminal state.
                If False, returns the prediction immediately. Defaults to False.
            wait_options (WaitOptions): Options for customizing wait behavior.
                If not provided, default options are used.

        Returns:
            PredictionWithIO: The new prediction.

        Raises:
            TypeError: When an invalid combination of version_id, model_owner,
                and model_name is provided.
        """
        path = self._prediction_path(model_owner, model_name, version_id)

        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=path,
                # 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

    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 (str): Unique identifier of the prediction.
            options (WaitOptions): Options for customizing wait behavior.
                If not provided, default options are used.

        Returns:
            PredictionWithIO: Prediction after completion.
        """
        while True:
            result = self.get(predictionID)
            if result["status"] in {Status.succeeded, Status.failed}:
                return result
            time.sleep(options.poll_interval)

    async def async_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 (str): Unique identifier of the prediction.
            options (WaitOptions): Options for customizing wait behavior.
                If not provided, default options are used.

        Returns:
            PredictionWithIO: Prediction after completion.
        """
        while True:
            result = await self.async_get(predictionID)
            if result["status"] in {Status.succeeded, Status.failed}:
                return result
            await asyncio.sleep(options.poll_interval)
