import logging
import os
from functools import partial
from pathlib import Path
from typing import List, Optional, Union
from uuid import UUID

import orjson
from beartype import beartype
from deprecation import deprecated

import picsellia.exceptions as exceptions
import picsellia.utils as utils
from picsellia.colors import Colors
from picsellia.decorators import exception_handler
from picsellia.sdk.connexion import Connexion
from picsellia.sdk.dao import Dao
from picsellia.sdk.datalake import Datalake
from picsellia.sdk.model_context import ModelContext
from picsellia.sdk.model_file import ModelFile
from picsellia.sdk.tag import Tag, TagTarget
from picsellia.sdk.taggable import Taggable
from picsellia.types.enums import Framework, InferenceType, ObjectDataType
from picsellia.types.schemas import ModelVersionSchema

logger = logging.getLogger("picsellia")


class ModelVersion(Dao, Taggable):
    def __init__(self, connexion: Connexion, data: dict):
        Dao.__init__(self, connexion, data)
        Taggable.__init__(self, TagTarget.MODEL_VERSION)

    @property
    def origin_name(self) -> str:
        """Name of origin of this (ModelVersion)"""
        return self._origin_name

    @property
    def origin_id(self) -> UUID:
        """UUID of the origin (Model) of this (ModelVersion)"""
        return self._origin_id

    @property
    def name(self) -> str:
        """Name of this (ModelVersion)"""
        return self._name

    @property
    def version(self) -> int:
        """Version number of this (ModelVersion)"""
        return self._version

    @property
    def type(self) -> InferenceType:
        """Type of this (ModelVersion)"""
        return self._type

    @property
    def framework(self) -> Framework:
        """Framework of this (ModelVersion)"""
        return self._framework

    @property
    def labels(self) -> Optional[dict]:
        return self._labels

    def __str__(self):
        return f"{Colors.BLUE}Version {self.name} of Model '{self.version}' with type {self.type.name} and framework {self.framework.name} {Colors.ENDC} (id: {self.id})"

    @exception_handler
    @beartype
    def get_resource_url_on_platform(self) -> str:
        """Get platform url of this resource.

        Examples:
            ```python
            print(foo_dataset.get_resource_url_on_platform())
            >>> "https://app.picsellia.com/model/62cffb84-b92c-450c-bc37-8c4dd4d0f590"
            ```

        Returns:
            Url on Platform for this resource
        """

        return f"{self.connexion.host}/model/version/{self.id}"

    @exception_handler
    @beartype
    def sync(self) -> dict:
        r = self.connexion.get(f"/sdk/model/version/{self.id}").json()
        self.refresh(r)
        return r

    @exception_handler
    @beartype
    def refresh(self, data: dict) -> ModelVersionSchema:
        schema = ModelVersionSchema(**data)
        self._origin_name = schema.origin.name
        self._origin_id = schema.origin.id
        self._version = schema.version
        self._name = schema.name
        self._type = schema.type
        self._framework = schema.framework
        self._labels = schema.labels
        return schema

    @exception_handler
    @beartype
    def get_tags(self) -> List[Tag]:
        """Retrieve the tags of your model version.

        Examples:
            ```python
            tags = my_model_version.get_tags()
            assert tags[0].name == "my-model-version-1"
            ```

        Returns:
            List of tags as Tag
        """
        r = self.sync()
        return list(map(partial(Tag, self.connexion), r["tags"]))

    @exception_handler
    @beartype
    def update(
        self,
        labels: Optional[dict] = None,
        docker_image_name: Optional[str] = None,
        docker_flags: Optional[List[str]] = None,
        thumb_object_name: Optional[str] = None,
        notebook_link: Optional[str] = None,
        base_parameters: Optional[dict] = None,
        docker_env_variables: Optional[dict] = None,
        framework: Union[str, Framework, None] = None,
        type: Union[str, InferenceType, None] = None,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> None:
        """Update this model version with some new infos.

        Examples:
            ```python
            model_v1.update(docker_image_name="docker.io/model1")
            ```
        """
        if framework:
            framework = Framework.validate(framework)

        if type:
            type = InferenceType.validate(type)

        if type and type not in [
            InferenceType.CLASSIFICATION,
            InferenceType.OBJECT_DETECTION,
            InferenceType.SEGMENTATION,
        ]:
            raise TypeError(f"Type '{type}' not supported yet for model version")

        payload = {
            "labels": labels,
            "docker_image_name": docker_image_name,
            "docker_flags": docker_flags,
            "thumb_object_name": thumb_object_name,
            "notebook_link": notebook_link,
            "base_parameters": base_parameters,
            "docker_env_variables": docker_env_variables,
            "type": type,
            "framework": framework,
            "name": name,
            "description": description,
        }
        filtered_payload = utils.filter_payload(payload)
        r = self.connexion.patch(
            f"/sdk/model/version/{self.id}",
            data=orjson.dumps(filtered_payload),
        ).json()
        self.refresh(r)

    @exception_handler
    @beartype
    def delete(self) -> None:
        """Delete model version.

        Delete the model in Picsellia database

        Examples:
            ```python
            model_v1.delete()
            ```
        """
        self.connexion.delete(f"/sdk/model/version/{self.id}")
        logger.info(f"{self} deleted.")

    @exception_handler
    @beartype
    def get_context(self) -> ModelContext:
        """Get ModelContext of this model

        Examples:
            ```python
            model_v1 = client.get_model(name="my-model").get_version(0)
            context = model_v1.get_context()
            context.get_infos()
            ```

        Returns:
            ModelContext objects that you can use and manipulate
        """
        try:
            r = self.connexion.get(f"/sdk/model/version/{self.id}/context").json()
        except exceptions.ForbiddenError:
            r = self.connexion.get(
                f"/sdk/model/version/{self.id}/public/context"
            ).json()

        return ModelContext(self.connexion, r)

    @exception_handler
    @beartype
    def list_files(self) -> List[ModelFile]:
        """Get a list of ModelFile that were stored with this model

        Examples:
            ```python
            model_v1 = client.get_model(name="my-model").get_version(0)
            files = model_v1.list_files()
            files[0].download()
            ```

        Returns:
            A list of (ModelFile) that you can use and manipulate
        """
        r = self.connexion.get(f"/sdk/model/version/{self.id}/modelfiles").json()
        return list(map(partial(ModelFile, self.connexion), r["items"]))

    @exception_handler
    @beartype
    def get_file(self, name: str) -> ModelFile:
        """Retrieve a ModelFile that were stored with this name into this model

        Examples:
            ```python
            model_v1 = client.get_model(name="my-model").get_version(0)
            file = model_v1.get_file("model-latest")
            file.download()
            ```

        Returns:
            A (ModelFile) that you can use and manipulate
        """
        params = {"name": name}
        r = self.connexion.get(
            f"/sdk/model/version/{self.id}/modelfiles/find", params=params
        ).json()
        return ModelFile(self.connexion, r)

    @exception_handler
    @beartype
    def store(
        self,
        name: str,
        path: Union[str, Path],
        do_zip: bool = False,
        replace: bool = False,
    ) -> ModelFile:
        """Store a file into picsellia storage and attach it to this model.

        Examples:
            ```python
            model.store("model-latest", "./lg_test_file.pb")
            ```
        Arguments:
            name (str): Name of file
            path (str or Path): Path of file to store
            do_zip (bool, optional): If true, zip directory to store it. Defaults to False.
            replace (bool, optional): If true, if a file with given name exists, it will be replaced. Defaults to False.

        Returns:
            A (ModelFile) object
        """
        if not os.path.exists(path):
            raise exceptions.FileNotFoundException(f"{path} not found")

        if do_zip:
            path = utils.zipdir(path)

        filename = os.path.split(path)[-1]
        object_name = self.connexion.generate_model_version_object_name(
            filename, ObjectDataType.MODEL_FILE, self.id
        )
        _, is_large, _ = self.connexion.upload_file(object_name, path)

        # Replacing file
        if replace:
            try:
                model_file = self.get_file(name)
                model_file.delete()
            except exceptions.ResourceNotFoundError:
                logger.debug(f"Model file `{name}` did not exist yet.")

        payload = {
            "name": name,
            "filename": filename,
            "object_name": object_name,
            "large": is_large,
        }
        r = self.connexion.post(
            f"/sdk/model/version/{self.id}/modelfiles",
            data=orjson.dumps(payload),
        ).json()
        model_file = ModelFile(self.connexion, r)
        logger.info(f"{model_file} stored for {self}")
        return model_file

    @exception_handler
    @beartype
    @deprecated(
        deprecated_in="6.7.0",
        details="update_thumbnail method will be removed in the future, you should use the web app instead",
    )
    def update_thumbnail(self, path: Union[str, Path]) -> None:
        """Updates the model thumbnail.

        Update the model thumbnail with given file.
        File size shall be less than 5Mb.

        Examples:
            ```python
            model.update_thumb("test.png")
            ```
        Arguments:
            path (str or Path): Path of the thumbnail you want to push

        Raises:
            FileNotFoundException: If there is no file in given path
            exceptions.InvalidQueryError: If file is too large
            exceptions.PicselliaError: If an unexpected error occurred while uploading file
        """
        if not os.path.isfile(path):
            raise exceptions.FileNotFoundException(f"{path} not found")

        filesize = Path(path).stat().st_size

        if filesize > 5 * 1024 * 1024:
            raise exceptions.BadRequestError("File too large limit is 5Mb")

        filename = os.path.split(path)[-1]
        object_name = self.connexion.generate_model_version_object_name(
            filename, ObjectDataType.MODEL_THUMB, self.id
        )
        self.connexion.upload_file(object_name, path)

        self.update(thumb_object_name=object_name)

    @exception_handler
    @beartype
    def deploy(
        self,
        name: Optional[str] = None,
        target_datalake: Optional[Datalake] = None,
        min_threshold: Optional[float] = None,
    ):
        """Create a (Deployment) for a model.

        This method allows you to create a (Deployment) on Picsellia. You will then have
        access to the monitoring dashboard and the model management part!

        Examples:
            ```python
            model_version = client.get_model(name="my-awesome-model").get_version(0)
            deployment = model_version.deploy(name="my-awesome-deployment", min_threshold=0.5)
            ```

        Arguments:
            name (str): Name of your deployment. Defaults to a random name.
            min_threshold (float): Threshold of detection scores used by models when predicting. Defaults to 0.
            target_datalake (Datalake): Datalake to use when data are pushed into Picsellia.
                Defaults to organization default Datalake.


        Returns:
            A (Deployment)
        """
        from picsellia.sdk.deployment import Deployment

        payload = {
            "name": name,
            "min_threshold": min_threshold,
        }
        if target_datalake is not None:
            payload["target_datalake_id"] = target_datalake.id
        filtered_payload = utils.filter_payload(payload)
        r = self.connexion.post(
            f"/sdk/model/version/{self.id}/deploy",
            data=orjson.dumps(filtered_payload),
        ).json()
        deployment = Deployment(self.connexion, r)
        logger.info(f"{deployment} created")
        return deployment
