"""Models module."""

from typing import Optional, TypedDict, cast

from wriftai._resource import Resource
from wriftai.common_types import (
    Model,
    ModelVersionWithDetails,
    ModelVisibility,
    NotRequired,
    SortDirection,
    StrEnum,
)
from wriftai.pagination import PaginatedResponse, PaginationOptions


class ModelWithDetails(Model):
    """Represents a model with details."""

    overview: str | None
    """The overview of the model."""
    latest_version: ModelVersionWithDetails | None
    """The details of the latest version of the model."""
    source_url: str | None
    """Source url from where the model's code can be referenced."""
    license_url: str | None
    """License url where the model's usage is specified."""
    paper_url: str | None
    """Paper url from where research info on the model
        can be found."""


class ModelsSortBy(StrEnum):
    """Enumeration of possible sorting options for querying models."""

    CREATED_AT = "created_at"
    PREDICTIONS_COUNT = "predictions_count"


class ModelPaginationOptions(PaginationOptions):
    """Pagination options for querying models."""

    sort_by: NotRequired[ModelsSortBy]
    """The sorting criteria."""

    sort_direction: NotRequired[SortDirection]
    """The sorting direction."""

    category_slugs: NotRequired[list[str]]
    """The list of category slugs to filter models."""


class UpdateModelParams(TypedDict):
    """Parameters for updating a model."""

    name: NotRequired[str]
    """The name of the model."""
    description: NotRequired[str | None]
    """Description of the model."""
    visibility: NotRequired[ModelVisibility]
    """The visibility of the model."""
    hardware_identifier: NotRequired[str]
    """The identifier of the hardware used by the model."""
    source_url: NotRequired[str | None]
    """Source url from where the model's code can be referenced."""
    license_url: NotRequired[str | None]
    """License url where the model's usage is specified."""
    paper_url: NotRequired[str | None]
    """Paper url from where research info on the model can be
        found."""
    overview: NotRequired[str | None]
    """The overview of the model."""
    category_slugs: NotRequired[list[str]]
    """List of model category slugs."""


class CreateModelParams(TypedDict):
    """Parameters for creating a model."""

    name: str
    """The name of the model."""
    hardware_identifier: str
    """The identifier of the hardware used by the model."""
    visibility: NotRequired[ModelVisibility]
    """The visibility of the model."""
    description: NotRequired[str | None]
    """Description of the model."""
    source_url: NotRequired[str | None]
    """Source url from where the model's code can be referenced."""
    license_url: NotRequired[str | None]
    """License url where the model's usage is specified."""
    paper_url: NotRequired[str | None]
    """Paper url from where research info on the model can be
        found."""
    overview: NotRequired[str | None]
    """The overview of the model."""
    category_slugs: NotRequired[list[str]]
    """List of model category slugs."""


class ModelsResource(Resource):
    """Resource for operations related to models."""

    _SEARCH_MODELS_SUFFIX = "/models"

    def delete(self, identifier: str) -> None:
        """Delete a model.

        Args:
            identifier: The model identifier in owner/name format (for example:
                deepseek-ai/deepseek-r1).
        """
        owner, name = self._parse_model_identifier(identifier=identifier)
        self._api.request("DELETE", f"{self._MODELS_API_PREFIX}/{owner}/{name}")

    async def async_delete(self, identifier: str) -> None:
        """Delete a model.

        Args:
            identifier: The model identifier in owner/name format (for example:
                deepseek-ai/deepseek-r1).
        """
        owner, name = self._parse_model_identifier(identifier=identifier)
        await self._api.async_request(
            "DELETE", f"{self._MODELS_API_PREFIX}/{owner}/{name}"
        )

    def list(
        self,
        pagination_options: Optional[ModelPaginationOptions] = None,
        owner: Optional[str] = None,
    ) -> PaginatedResponse[Model]:
        """List models.

        Args:
            pagination_options: Optional settings to control pagination behavior.
            owner: Username of the model's owner to fetch models for. If None, all
                models are fetched.

        Returns:
            Paginated response containing models and navigation metadata.
        """
        path = self._build_list_path(owner)

        response = self._api.request(method="GET", params=pagination_options, path=path)

        # The response will always match the ModelPaginatedResponse 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[ModelPaginationOptions] = None,
        owner: Optional[str] = None,
    ) -> PaginatedResponse[Model]:
        """List models.

        Args:
            pagination_options: Optional settings to control pagination behavior.
            owner: Username of the model's owner to fetch models for. If None, all
                models are fetched.

        Returns:
            Paginated response containing models and navigation metadata.
        """
        path = self._build_list_path(owner)

        response = await self._api.async_request(
            method="GET", params=pagination_options, path=path
        )
        # The response will always match the ModelPaginatedResponse 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 _build_list_path(self, owner: Optional[str] = None) -> str:
        """Construct the API path for listing models.

        Args:
            owner: Username of the model's owner to fetch models for. If None, returns
                a path to fetch all models.

        Returns:
            The constructed API path for listing models.
        """
        return (
            f"{self._MODELS_API_PREFIX}/{owner}" if owner else self._MODELS_API_PREFIX
        )

    def get(self, identifier: str) -> ModelWithDetails:
        """Get a model.

        Args:
            identifier: The model identifier in owner/name format (for example:
                deepseek-ai/deepseek-r1).

        Returns:
            A model object.
        """
        owner, name = self._parse_model_identifier(identifier=identifier)
        response = self._api.request("GET", f"{self._MODELS_API_PREFIX}/{owner}/{name}")

        return cast(ModelWithDetails, response)

    async def async_get(self, identifier: str) -> ModelWithDetails:
        """Get a model.

        Args:
            identifier: The model identifier in owner/name format (for example:
                deepseek-ai/deepseek-r1).

        Returns:
            A model object.
        """
        owner, name = self._parse_model_identifier(identifier=identifier)
        response = await self._api.async_request(
            "GET", f"{self._MODELS_API_PREFIX}/{owner}/{name}"
        )

        return cast(ModelWithDetails, response)

    def create(self, params: CreateModelParams) -> ModelWithDetails:
        """Create a model.

        Args:
            params: Model creation parameters.

        Returns:
            The new model.
        """
        response = self._api.request(
            method="POST",
            path=self._MODELS_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]
        )

        return cast(ModelWithDetails, response)

    async def async_create(self, params: CreateModelParams) -> ModelWithDetails:
        """Create a model.

        Args:
            params: Model creation parameters.

        Returns:
            The new model.
        """
        response = await self._api.async_request(
            method="POST",
            path=self._MODELS_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]
        )

        return cast(ModelWithDetails, response)

    def update(
        self,
        identifier: str,
        params: UpdateModelParams,
    ) -> ModelWithDetails:
        """Update a model.

        Args:
            identifier: The model identifier in owner/name format (for example:
                deepseek-ai/deepseek-r1).
            params: The fields to update.

        Returns:
            The updated model.
        """
        owner, name = self._parse_model_identifier(identifier=identifier)
        response = self._api.request(
            method="PATCH",
            path=f"{self._MODELS_API_PREFIX}/{owner}/{name}",
            # 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]
        )
        return cast(ModelWithDetails, response)

    async def async_update(
        self,
        identifier: str,
        params: UpdateModelParams,
    ) -> ModelWithDetails:
        """Update a model.

        Args:
            identifier: The model identifier in owner/name format (for example:
                deepseek-ai/deepseek-r1).
            params: The fields to update.

        Returns:
            The updated model.
        """
        owner, name = self._parse_model_identifier(identifier=identifier)
        response = await self._api.async_request(
            method="PATCH",
            path=f"{self._MODELS_API_PREFIX}/{owner}/{name}",
            # 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]
        )

        return cast(ModelWithDetails, response)

    def search(
        self,
        q: str,
        pagination_options: Optional[PaginationOptions] = None,
    ) -> PaginatedResponse[Model]:
        """Search models.

        Args:
            q: The search query.
            pagination_options: Optional settings to control pagination behavior.

        Returns:
            Paginated response containing models and navigation metadata.
        """
        params = self._build_search_params(q, pagination_options)

        response = self._api.request(
            method="GET",
            params=params,
            path=f"{self._SEARCH_API_PREFIX}{self._SEARCH_MODELS_SUFFIX}",
        )
        # 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_search(
        self,
        q: str,
        pagination_options: Optional[PaginationOptions] = None,
    ) -> PaginatedResponse[Model]:
        """Search models.

        Args:
            q: The search query.
            pagination_options: Optional settings to control pagintation behavior.

        Returns:
            Paginated response containing models and navigation metadata.
        """
        params = self._build_search_params(q, pagination_options)

        response = await self._api.async_request(
            method="GET",
            params=params,
            path=f"{self._SEARCH_API_PREFIX}{self._SEARCH_MODELS_SUFFIX}",
        )
        # 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]
