from __future__ import annotations

import logging
import pickle  # nosec
from typing import TYPE_CHECKING, Any, List

from vectice.models.attachment import TAttachment, TFormattedAttachment

if TYPE_CHECKING:
    from mlflow import MlflowClient  # pyright: ignore[reportPrivateImportUsage]

from vectice.models import AdditionalInfo, ExtraInfo, Metric, Property
from vectice.models.dataset import TDerivedFrom, get_derived_from
from vectice.utils.common_utils import check_key_duplicates, format_attachments, format_metrics, format_properties

_logger = logging.getLogger(__name__)


class Model:
    """Represent a wrapped model.

    To document models within an iteration, you need to wrap it as a Vectice model. It involves capturing metadata about the predictor.
    This metadata enables Vectice to provide summaries and insights about the model.
    """

    def __init__(
        self,
        library: str | None = None,
        technique: str | None = None,
        metrics: dict[str, int | float] | list[Metric] | Metric | None = None,
        properties: dict[str, str | int] | list[Property] | Property | None = None,
        name: str | None = None,
        attachments: TAttachment | None = None,
        predictor: Any = None,
        derived_from: list[TDerivedFrom] | None = None,
        additional_info: dict[str, str] | AdditionalInfo | None = None,
    ):
        """Wrap a model (predictor).

        A Vectice Model is a wrapped predictor that can be logged
        inside a Vectice iteration.

        Parameters:
            library: The library used to generate the model.
            technique: The modeling technique used.
            metrics: A dict for example `{"MSE": 1}`.
            properties: A dict, for example `{"folds": 32}`.
            name: The model name. If None, will be auto-generated by Vectice.
            attachments: The file paths that will be attached to the iteration along with the predictor.
            predictor: The predictor.
            derived_from: List of dataset versions or model versions (or their ids) to link as lineage.
            additional_info: An optional set of key values related to the context in which the model was created. For instance if an experiment tracker was used or a specific framework.
        """
        from vectice import pickle_capture

        self._library = library
        self._technique = technique
        self._name = name or self._generate_name()
        self._metrics = format_metrics(metrics) if metrics else None
        self._properties = format_properties(properties) if properties else None
        self._attachments = format_attachments(attachments) if attachments else None
        self._derived_from = get_derived_from(derived_from)
        self._additional_info = self._format_additional_info(additional_info) if additional_info else None
        self._latest_version_id: str | None = None
        self._predictor = predictor if pickle_capture else None  # nosec

    def __repr__(self):
        return (
            f"Model(name='{self.name or ''}', library='{self.library or ''}', technique='{self.technique or ''}', "
            f"metrics={self.metrics}, properties={self.properties}, attachments={self.attachments}, additional_info={self.additional_info})"
        )

    @property
    def latest_version_id(self) -> str | None:
        """The id of the latest version of this model.

        Returns:
            The id of the latest version of this model.
        """
        return self._latest_version_id

    @latest_version_id.setter
    def latest_version_id(self, value: str) -> None:
        """Set the id of the latest version of this model.

        Parameters:
            value: The id of the latest version of this model.
        """
        self._latest_version_id = value

    @property
    def name(self) -> str | None:
        """The model's name.

        Returns:
            The model's name.
        """
        return self._name

    @name.setter
    def name(self, name: str):
        """Set the model's name.

        Parameters:
            name: The name of the model.
        """
        self._name = name

    @property
    def predictor(self) -> Any:
        """The model's predictor.

        Returns:
            The model's predictor.
        """
        return self._predictor  # noqa: S301

    @predictor.setter
    def predictor(self, predictor: Any):
        """Set the model's predictor.

        Parameters:
            predictor: The predictor of the model.
        """
        self._predictor = predictor

    @classmethod  # pyright: ignore[reportArgumentType]
    def read_predictor_file(cls, path: str) -> Any:
        with open(path, "rb") as fh:
            return pickle.load(fh)  # noqa: S301

    @property
    def library(self) -> str | None:
        """The name of the library used to generate the model.

        Returns:
            The name of the library used to generate the model.
        """
        return self._library

    @library.setter
    def library(self, library: str):
        """Set the name of the library used to create the model.

        Parameters:
            library: The name of the library used to create the model.
        """
        self._library = library

    @property
    def technique(self) -> str | None:
        """The name of the modeling technique used to learn the model.

        Returns:
            The name of the modeling technique used to learn the model.
        """
        return self._technique

    @technique.setter
    def technique(self, technique: str):
        """Set the name of the modeling technique used to learn the model.

        Parameters:
            technique: The modeling technique used.
        """
        self._technique = technique

    @property
    def metrics(self) -> list[Metric] | None:
        """The model's metrics.

        Returns:
            The model's metrics.
        """
        return self._metrics

    @metrics.setter
    def metrics(self, metrics: dict[str, int | float] | list[Metric] | Metric | None):
        """Set the model's metrics.

        Parameters:
            metrics: The metrics of the model.
        """
        _logger.warning("To save your updated model metrics, you must reassign your model to an iteration.")
        self._metrics = format_metrics(metrics)

    @property
    def properties(self) -> list[Property] | None:
        """The model's properties.

        Returns:
            The model's properties.
        """
        return self._properties

    @properties.setter
    def properties(self, properties: dict[str, str | int] | list[Property] | Property | None):
        """Set the model's properties.

        Parameters:
            properties: The properties of the model.
        """
        _logger.warning("To save your updated model properties, you must reassign your model to an iteration.")
        self._properties = format_properties(properties) if properties else None

    @property
    def attachments(self) -> list[TFormattedAttachment] | None:
        """The file attachment objects or paths associated with the model.

        Returns:
        `List[Union[Table, str]] | None`
            A list of attachment objects, where each attachment is one of:

        - `Table`
            A formatted table object containing structured data
        - `str`
            A string representation of an attachment (e.g., file path)
        - `None`
            If no attachments are present
        """
        return self._attachments

    @attachments.setter
    def attachments(self, attachments: TAttachment):
        """Attach a file or files to the model.

        Parameters:
            attachments: The filename or filenames of the file or set of files to attach to the model.
        """
        self._attachments = format_attachments(attachments)

    @property
    def additional_info(self) -> AdditionalInfo | None:
        """The additional info associated with the model.

        Returns:
            The additional info associated with the model.
        """
        return self._additional_info

    @additional_info.setter
    def additional_info(self, additional_info: dict[str, str] | AdditionalInfo):
        """Attach an additional info to the model.

        Parameters:
            additional_info: The additional information to be passed to the model linked to Experiment Trackers or other context in which the model was created.
        """
        self._additional_info = self._format_additional_info(additional_info)

    @property
    def derived_from(self) -> list[str]:
        """The datasets versions or models versions from which this model is derived.

        Returns:
            The datasets versions or models versions from which this model is derived.
        """
        return self._derived_from

    def _generate_name(self) -> str | None:
        if self.library is not None and self.technique is not None:
            return f"{self.library} {self.technique} model"
        return None

    @staticmethod  # pyright: ignore[reportArgumentType]
    def _remove_incorrect_extra_info(extra_info: list[ExtraInfo]) -> list[ExtraInfo]:
        for info in extra_info:
            if not isinstance(info, ExtraInfo):  # pyright: ignore[reportUnnecessaryIsInstance]
                logging.warning(f"Incorrect extra info '{info}'. Please check extra info type.")
                extra_info.remove(info)
        return extra_info

    def _format_additional_info(self, additional_info: dict[str, str] | AdditionalInfo | None) -> AdditionalInfo | None:
        if additional_info is None:
            return None
        if isinstance(additional_info, AdditionalInfo):
            if isinstance(additional_info.extra_info, list):
                extra_info = self._remove_incorrect_extra_info(additional_info.extra_info)
                additional_info.extra_info = extra_info
                key_list = [info.key for info in additional_info.extra_info]
                check_key_duplicates(key_list)
            return additional_info

        if isinstance(additional_info, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
            extra_info_list: List[ExtraInfo] = []
            for key, value in additional_info.items():
                extra_info_list.append(ExtraInfo(key=key, value=value))
            return AdditionalInfo(url=None, run=None, framework=None, extra_info=extra_info_list)
        else:
            raise ValueError("Please check additional info type.")

    @staticmethod  # pyright: ignore[reportArgumentType]
    def mlflow(
        run_id: str,
        client: MlflowClient,
        url: str | None = None,
        derived_from: list[TDerivedFrom] | None = None,
    ) -> Model:
        """Extract automatically information from an Mlflow run into a Vectice model to be assigned to an iteration.

        ```python
        import mlflow
        from mlflow.client import MlflowClient
        from vectice import Model

        model = Model.mlflow(
                run_id="d479ca87954f4a0abc4da3333b990cb3",
                client=MlflowClient() #also work with client=mlflow
            )
        ```


        Parameters:
            run_id: The run id of the mlflow model to be wrapped.
            client: The mlflow client in order to access the run through get_run() method. WARNING: the client should have the tracking URI where the run is located.
            url: The Mlflow UI URL to access the run info.
            derived_from: List of dataset versions or model versions (or their ids) to link as lineage.
        """
        from vectice.models.model_mlflow import MlflowModel

        model = MlflowModel(run_id, client=client, url=url, derived_from=derived_from).as_vectice_model()
        return model
