# Copyright (C) 2023 - 2026 ANSYS, Inc. and/or its affiliates.
# SPDX-License-Identifier: MIT
#
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import logging
from typing import Any, BinaryIO, List, Optional, Tuple, Union

from pydantic import BaseModel, PositiveInt, ValidationError

from ansys.simai.core.data.base import ComputableDataModel, Directory
from ansys.simai.core.data.geomai.workspaces import GeomAIWorkspace
from ansys.simai.core.data.types import (
    File,
    Identifiable,
    get_id_from_identifiable,
)
from ansys.simai.core.errors import InvalidArguments

logger = logging.getLogger(__name__)


class GeomAIPredictionConfiguration(BaseModel):
    """The configuration used to run a GeomAI prediction."""

    latent_params: List[float]
    """A list of floats that represent the position of the geometry in the latent space.

    These parameters describe the shape in a compressed form.
    The number of floats should match the ``nb_latent_param`` your model was requested with.

    Required.
    """
    resolution: Optional[Tuple[PositiveInt, PositiveInt, PositiveInt]] = None
    """A list of three integers defining the number of voxels along the X, Y, and Z axes.

    Use higher resolution for complex or precise geometries, and lower resolution for simple shapes or quick previews.

    The total number of voxels must not exceed 900^3, that is `x`, `y`, `z` multiplied together must be less than or equal to 900^3.
    If you exceed that value, an error will occur.

    Defaults to ``[100,100,100]``, if ``None`` is provided.

    For the maximum resolution of 900^3, the prediction takes approximately 10 minutes (approximately 1 microsecond per voxel).
    """

    def __init__(self, *args, **kwargs):
        """Raises :exc:`~ansys.simai.core.errors.InvalidArguments` if the input data cannot be validated to from a valid model."""
        try:
            super().__init__(*args, **kwargs)
        except ValidationError as e:
            raise InvalidArguments(e.errors(include_url=False)) from None


class GeomAIPrediction(ComputableDataModel):
    """Provides the local representation of a GeomAI prediction object."""

    @property
    def configuration(self) -> GeomAIPredictionConfiguration:
        """The configuration used to run the prediction."""
        return GeomAIPredictionConfiguration.model_construct(**self.fields["configuration"])

    def delete(self) -> None:
        """Remove a prediction from the server."""
        self._client._api.delete_geomai_prediction(self.id)
        self._unregister()

    def download(self, file: Optional[File] = None) -> Union[None, BinaryIO]:
        """Download the file generated by the prediction.

        The downloaded file will be in the VTP format.

        Args:
            file: Binary file-object or the path of the file to put the content into.

        Returns:
            ``None`` if a file is specified or a binary file-object otherwise.
        """
        self.wait()
        return self._client._api.download_geomai_prediction(self.id, file)


class GeomAIPredictionDirectory(Directory[GeomAIPrediction]):
    """Provides a collection of methods related to GeomAI model predictions.

    This method is accessed through ``client.geomai.predictions``.

    Example:
        .. code-block:: python

            import ansys.simai.core as asc

            simai_client = asc.from_config()
            simai_client.geomai.predictions.list()

    """

    _data_model = GeomAIPrediction

    def list(
        self, workspace: Optional[Identifiable[GeomAIWorkspace]] = None
    ) -> List[GeomAIPrediction]:
        """List all GeomAI predictions on the server that belong to the specified workspace or the configured one.

        Args:
            workspace: ID or :class:`model <.workspaces.GeomAIWorkspace>` of the workspace to list the predictions for.
                This parameter is necessary if no workspace is set for the client.
        """
        workspace_id = get_id_from_identifiable(
            workspace, default=self._client.geomai._current_workspace
        )
        raw_predictions = self._client._api.geomai_predictions(workspace_id)
        return list(map(self._model_from, raw_predictions))

    def get(self, id: str) -> GeomAIPrediction:
        """Get a specific GeomAI prediction object from the server by ID.

        Args:
            id: ID of the prediction.

        Returns:
            :class:`GeomAIPrediction` instance with the given ID if it exists.

        Raises:
            NotFoundError: No prediction with the given ID exists.
        """
        return self._model_from(self._client._api.get_geomai_prediction(id))

    def delete(self, prediction: Identifiable[GeomAIPrediction]) -> None:
        """Delete a specific GeomAI prediction from the server.

        Args:
            prediction: ID or :class:`model <GeomAIPrediction>` of the prediction.

        Raises:
            NotFoundError: No prediction with the given ID exists.
        """
        prediction_id = get_id_from_identifiable(prediction)
        self._client._api.delete_geomai_prediction(prediction_id)
        self._unregister_item_with_id(prediction_id)

    def run(  # noqa: D417
        self,
        configuration: Union[GeomAIPredictionConfiguration, dict[str, Any]],
        workspace: Optional[Identifiable[GeomAIWorkspace]] = None,
    ) -> GeomAIPrediction:
        """Run a prediction in the given workspace with the given configuration.

        Args:
            configuration: The configuration to run the prediction with.
            workspace: Optional ID or :class:`model <.workspaces.GeomAIWorkspace>` of the target workspace.
                Defaults to the current workspace if set.

        Returns:
            Created prediction object.

        Raises:
            ProcessingError: If the server failed to process the request.

        Examples:
            .. code-block:: python

                import ansys.simai.core as asc

                simai_client = asc.from_config()
                workspace = simai_client.geomai.workspaces.list()[0]
                prediction = simai_client.geomai.predictions.run(
                    dict(latent_params=[0.1, 1.2, 0.76], resolution=(100, 100, 100)),
                    workspace,
                )

        """
        if not isinstance(configuration, GeomAIPredictionConfiguration):
            try:
                configuration = GeomAIPredictionConfiguration.model_validate(configuration)
            except ValidationError as e:
                raise InvalidArguments(e.errors(include_url=False)) from None
        workspace_id = get_id_from_identifiable(
            workspace, default=self._client.geomai._current_workspace
        )
        return self._model_from(
            self._client._api.run_geomai_prediction(
                workspace_id, configuration.model_dump(exclude_unset=True)
            )
        )

    def download(
        self, prediction: Identifiable[GeomAIPrediction], file: Optional[File] = None
    ) -> Union[None, BinaryIO]:
        """Download the file generated by the prediction.

        The downloaded file will be in the VTP format.

        Args:
            prediction: ID or :class:`model <GeomAIPrediction>` of the prediction.
            file: Binary file-object or the path of the file to put the content into.

        Returns:
            ``None`` if a file is specified or a binary file-object otherwise.

        Note:
            Consider using :meth:`GeomAIPrediction.download` instead which will wait for the
                prediction to be ready instead of throwing an error if it's not.
        """
        prediction_id = get_id_from_identifiable(prediction)
        return self._client._api.download_geomai_prediction(prediction_id, file)
