# Copyright (C) 2025 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, Field, 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 divisions along the X, Y, and Z axes.

    Higher values allow for more detailed geometry reconstruction, while lower values reduce detail but speed up computation.

    Defaults to ``[100, 100, 100]`` if none is provided.
    """
    margin: Optional[float] = Field(default=None, ge=0, le=1)
    """A float that sets the size of the isosurface for the reconstruction of the geometry.

    Note:
        A margin of 0.0 is highly recommended for non-expert users.

    If you are an expert user, and if the generated geometry is noisy,
    you can try to adjust both resolution and margin values to find
    the right balance between them to generate a smoother geometry.

    | A higher margin gives a coarser surface with less detail.
    | A lower margin produces a sharper surface.

    Defaults to 0 if none is provided.
    """

    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.

        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.
        """
        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

            simai = ansys.simai.core.from_config()
            simai.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

                simai = ansys.simai.core.from_config()
                workspace = simai.geomai.workspaces.list()[0]
                prediction = simai.geomai.predictions.run(
                    dict(latent_params=[0.1, 1.2, 0.76], resolution=(100, 100, 100), margin=0.0),
                    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())
        )

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

        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.
        """
        prediction_id = get_id_from_identifiable(prediction)
        return self._client._api.download_geomai_prediction(prediction_id, file)
