"""Prediction-related types mirroring the Skytells API response shapes."""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional


class PredictionStatus(str, Enum):
    """
    Lifecycle status of a prediction.

    - ``pending``    — Queued, waiting to start.
    - ``starting``   — Allocating resources.
    - ``started``    — Resources allocated, about to process.
    - ``processing`` — Actively running.
    - ``succeeded``  — Completed successfully (output available).
    - ``failed``     — Completed with an error.
    - ``cancelled``  — Cancelled by the user.
    """

    PENDING = "pending"
    STARTING = "starting"
    STARTED = "started"
    PROCESSING = "processing"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    CANCELLED = "cancelled"


class PredictionType(str, Enum):
    """The type of prediction workload."""

    INFERENCE = "inference"
    TRAINING = "training"


class PredictionSource(str, Enum):
    """The source that created the prediction."""

    API = "api"
    CLI = "cli"
    WEB = "web"


class PredictionRequest:
    """
    Payload for creating a prediction via the Skytells API.

    Used by :meth:`~skytells.SkytellsClient.predict` and
    :meth:`~skytells.PredictionsAPI.create`.

    Args:
        model:   Model slug to run (e.g. ``"flux-pro"``).
        input:   Key-value input parameters for the model.
        webhook: Optional webhook configuration dict with ``url`` and ``events``.
        await_:  If ``True``, blocks until completion (default: ``False``).
        stream:  If ``True``, enables streaming (default: ``False``).

    Example::

        req = PredictionRequest(
            model="flux-pro",
            input={"prompt": "A sunset"},
            await_=True,
        )
    """

    def __init__(
        self,
        model: str,
        input: Dict[str, Any],
        webhook: Optional[Dict[str, Any]] = None,
        await_: bool = False,
        stream: bool = False,
    ) -> None:
        self.model = model
        self.input = input
        self.webhook = webhook
        self.await_ = await_
        self.stream = stream

    def to_dict(self) -> Dict[str, Any]:
        d: Dict[str, Any] = {
            "model": self.model,
            "input": self.input,
        }
        if self.webhook is not None:
            d["webhook"] = self.webhook
        if self.await_:
            d["await"] = True
        if self.stream:
            d["stream"] = True
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "PredictionRequest":
        return cls(
            model=d["model"],
            input=d.get("input", {}),
            webhook=d.get("webhook"),
            await_=d.get("await", False),
            stream=d.get("stream", False),
        )


class RunOptions:
    """
    Options for :meth:`~skytells.SkytellsClient.run`.

    Same as :class:`PredictionRequest` but without ``model`` (passed separately)
    and without ``await_`` (handled internally).

    Args:
        input:   Key-value input parameters for the model.
        webhook: Optional webhook configuration.
        stream:  Enable streaming (default: ``False``).
    """

    def __init__(
        self,
        input: Dict[str, Any],
        webhook: Optional[Dict[str, Any]] = None,
        stream: bool = False,
    ) -> None:
        self.input = input
        self.webhook = webhook
        self.stream = stream

    def to_dict(self) -> Dict[str, Any]:
        d: Dict[str, Any] = {"input": self.input}
        if self.webhook is not None:
            d["webhook"] = self.webhook
        if self.stream:
            d["stream"] = True
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "RunOptions":
        return cls(
            input=d.get("input", {}),
            webhook=d.get("webhook"),
            stream=d.get("stream", False),
        )


class WaitOptions:
    """
    Options for :meth:`~skytells.SkytellsClient.wait` to control polling.

    Args:
        interval: Polling interval in milliseconds (default: ``5000``).
        max_wait: Maximum time to wait in milliseconds.
                  Raises :class:`~skytells.SkytellsError` with ``WAIT_TIMEOUT`` if exceeded.
    """

    def __init__(
        self,
        interval: int = 5000,
        max_wait: Optional[int] = None,
    ) -> None:
        self.interval = interval
        self.max_wait = max_wait

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "WaitOptions":
        return cls(
            interval=d.get("interval", 5000),
            max_wait=d.get("max_wait"),
        )


class PredictionsListOptions:
    """
    Options for listing predictions with filters and pagination.

    Args:
        page:  Page number (default: ``1``).
        since: Include predictions on or after this date (``YYYY-MM-DD``).
        until: Include predictions on or before this date (``YYYY-MM-DD``).
        model: Filter by model slug.
    """

    def __init__(
        self,
        page: Optional[int] = None,
        since: Optional[str] = None,
        until: Optional[str] = None,
        model: Optional[str] = None,
    ) -> None:
        self.page = page
        self.since = since
        self.until = until
        self.model = model

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "PredictionsListOptions":
        return cls(
            page=d.get("page"),
            since=d.get("since"),
            until=d.get("until"),
            model=d.get("model"),
        )


class QueueItem:
    """An item in the local prediction queue."""

    def __init__(self, request: PredictionRequest) -> None:
        self.request = request
