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

from typing import TYPE_CHECKING, Dict, List, Optional

from ansys.simai.core.data.base import ComputableDataModel, Directory
from ansys.simai.core.data.types import (
    Filters,
    Identifiable,
    MonitorCallback,
    NamedFile,
    SizedIterator,
    get_id_from_identifiable,
    to_raw_filters,
    unpack_named_file,
)
from ansys.simai.core.errors import InvalidArguments
from ansys.simai.core.utils.pagination import DataModelIterator

if TYPE_CHECKING:
    from ansys.simai.core.data.geomai.projects import GeomAIProject
    from ansys.simai.core.data.geomai.training_data_parts import GeomAITrainingDataPart


def _upload_geomai_training_data_part(id, named_part, client, monitor_callback):
    with unpack_named_file(named_part) as (file, name, extension):
        (training_data_part_fields, upload_id) = client._api.create_geomai_training_data_part(
            id, name, extension
        )
        training_data_part = client.geomai.training_data_parts._model_from(
            training_data_part_fields, is_upload_complete=False
        )
        parts = client._api.upload_parts(
            f"geomai/training-data-parts/{training_data_part.id}/part",
            file,
            upload_id,
            monitor_callback=monitor_callback,
        )
        client._api.complete_geomai_training_data_part_upload(
            training_data_part.id, upload_id, parts
        )
    return training_data_part


class GeomAITrainingData(ComputableDataModel):
    """Provides the local representation of a training data object."""

    def __repr__(self) -> str:
        return f"<GeomAITrainingData: {self.id}, {self.name}>"

    @property
    def name(self) -> str:
        """Name of the training data."""
        return self.fields["name"]

    @property
    def parts(self) -> List["GeomAITrainingDataPart"]:
        """List of all :class:`parts<ansys.simai.core.data.geomai.training_data_parts.GeomAITrainingDataPart>`
        objects in the training data.
        """
        return [
            self._client.geomai.training_data_parts._model_from(training_data_part)
            for training_data_part in self.fields["parts"]
        ]

    @property
    def extracted_metadata(self) -> Optional[Dict]:
        """Metadata extracted from the training data."""
        return self.fields["extracted_metadata"]

    def extract_data(self) -> None:
        """Extract or reextract the data from a training data.

        Data should be extracted from a training data once all its parts have been fully uploaded.
        This is done automatically when using :meth:`GeomAITrainingDataDirectory.create_from_file` to create training data.

        Data can only be reextracted from a training data if the extraction previously failed or if new files have been added.
        """
        self._client._api.compute_geomai_training_data(self.id)

    def delete(self) -> None:
        """Delete the training data on the server."""
        self._client._api.delete_geomai_training_data(self.id)

    def upload_part(
        self, file: NamedFile, monitor_callback: Optional[MonitorCallback] = None
    ) -> "GeomAITrainingDataPart":
        """Add a part to the training data.

        Args:
            file: :obj:`~ansys.simai.core.data.types.NamedFile` to upload.
            monitor_callback: Optional callback for monitoring the progress of the download.
                For more information, see the :obj:`~ansys.simai.core.data.types.MonitorCallback`
                object.

        Returns:
            Created :class:`~ansys.simai.core.data.geomai.training_data_parts.GeomAITrainingDataPart`.
        """
        return _upload_geomai_training_data_part(self.id, file, self._client, monitor_callback)

    def add_to_project(self, project: Identifiable["GeomAIProject"]):
        """Add the training data to a :class:`~.projects.GeomAIProject` object.

        Args:
            project: ID or :class:`model <.projects.GeomAIProject>` object of the project to add the data to.
        """
        project_id = get_id_from_identifiable(project)
        self._client._api.add_geomai_training_data_to_project(self.id, project_id)

    def remove_from_project(self, project: Identifiable["GeomAIProject"]):
        """Remove the training data from a :class:`~.projects.GeomAIProject` object.

        Args:
            project: ID or :class:`model <.projects.GeomAIProject>` of the project to remove data from.

        Raises:
            ansys.simai.core.errors.ApiClientError: If the data is the project's sample.
            ansys.simai.core.errors.ApiClientError: If the project is in training.
        """
        project_id = get_id_from_identifiable(project)
        self._client._api.remove_geomai_training_data_from_project(self.id, project_id)


class GeomAITrainingDataDirectory(Directory[GeomAITrainingData]):
    """Provides a collection of methods related to training data.

    This class is accessed through ``client.training_data``.

    Example:
        List all of the GeomAI training data::

            import ansys.simai.core as asc

            simai_client = asc.from_config()
            simai_client.geomai.training_data.list()
    """

    _data_model = GeomAITrainingData

    def iter(self, filters: Optional[Filters] = None) -> SizedIterator[GeomAITrainingData]:
        """Iterate over all :class:`GeomAITrainingData` objects on the server.

        Args:
            filters: Optional :obj:`~ansys.simai.core.data.types.Filters` to apply.

        Returns:
            Iterator over all :class:`GeomAITrainingData` objects on the server.
        """
        raw_filters = to_raw_filters(filters)
        raw_iterable = self._client._api.iter_geomai_training_data(filters=raw_filters)
        return DataModelIterator(raw_iterable, self)

    def list(self, filters: Optional[Filters] = None) -> List[GeomAITrainingData]:
        """List all :class:`GeomAITrainingData` objects on the server.

        Warning:
            This can take a very long time, consider using :py:meth:`~iter` instead.

        Args:
            filters: Optional :obj:`~ansys.simai.core.data.types.Filters` to apply.

        Returns:
            List of all :class:`GeomAITrainingData` objects on the server.
        """
        return list(self.iter(filters))

    def get(self, id: Optional[str] = None, name: Optional[str] = None) -> GeomAITrainingData:
        """Get a specific :class:`GeomAITrainingData` object from the server.

        You can specify either the ID or the name, not both.

        Args:
            id: ID of the training data.
            name: Name of the training data.

        Returns:
            :class:`GeomAITrainingData`

        Raises:
            NotFoundError: No training data with the given ID exists.
        """
        if name and id:
            raise InvalidArguments("Cannot specify both 'id' and 'name' arguments.")
        elif name:
            return self._model_from(self._client._api.get_geomai_training_data_by_name(name))
        elif id:
            return self._model_from(self._client._api.get_geomai_training_data(id))
        else:
            raise InvalidArguments("Either 'id' or 'name' argument must be specified.")

    def delete(self, training_data: Identifiable[GeomAITrainingData]) -> None:
        """Delete a :class:`GeomAITrainingData` object and its associated parts from the server.

        Args:
            training_data: ID or :class:`model <GeomAITrainingData>` object of the :class:`GeomAITrainingData` object.
        """
        return self._client._api.delete_geomai_training_data(
            get_id_from_identifiable(training_data)
        )

    def create(
        self, name: str, project: Optional[Identifiable["GeomAIProject"]] = None
    ) -> GeomAITrainingData:
        """Create an empty :class:`GeomAITrainingData` object.

        Args:
            name: Name to give the new :class:`GeomAITrainingData` object.
            project: :class:`~.projects.Project` object to associate the data with.

        Returns:
            Created :class:`GeomAITrainingData` object.

        See Also:
            :meth:`create_from_file` handles all the creation of training data and parts for you
        """
        project_id = get_id_from_identifiable(project, required=False)
        return self._model_from(self._client._api.create_geomai_training_data(name, project_id))

    def create_from_file(
        self,
        file: NamedFile,
        project: Optional[Identifiable["GeomAIProject"]] = None,
        monitor_callback: Optional[MonitorCallback] = None,
    ) -> GeomAITrainingData:
        """Convenience function to create a :class:`GeomAITrainingData` object from a file.

        Args:
            file: :obj:`~ansys.simai.core.data.types.NamedFile` to upload.
            project: Optional project to attach the training data to.
                Defaults to the configured project if any.
            monitor_callback: Optional callback for monitoring the progress of the download.
                For more information, see the :obj:`~ansys.simai.core.data.types.MonitorCallback`
                object.
        """
        project_id = get_id_from_identifiable(
            project, default=self._client.geomai._current_project, required=False
        )
        with unpack_named_file(file) as (file_bytes, name, extension):
            training_data = self._model_from(
                self._client._api.create_geomai_training_data(name, project_id)
            )
            _upload_geomai_training_data_part(
                training_data.id, file, self._client, monitor_callback
            )
            self._client._api.compute_geomai_training_data(training_data.id)
            return training_data

    def upload_part(
        self,
        training_data: Identifiable[GeomAITrainingData],
        file: NamedFile,
        monitor_callback: Optional[MonitorCallback] = None,
    ) -> "GeomAITrainingDataPart":
        """Add a part to a :class:`GeomAITrainingData` object.

        Args:
            training_data: ID or :class:`model <GeomAITrainingData>` object of the training data to
                add the part to.
            file: :obj:`~ansys.simai.core.data.types.NamedFile` to upload.
            monitor_callback: Optional callback for monitoring the progress of the upload.
                For more information, see the :obj:`~ansys.simai.core.data.types.MonitorCallback`
                object.

        Returns:
            Added :class:`~ansys.simai.core.data.geomai.training_data_parts.GeomAITrainingDataPart` object.
        """
        return _upload_geomai_training_data_part(
            get_id_from_identifiable(training_data),
            file,
            self._client,
            monitor_callback,
        )
