from enum import Enum
from .wallaroo_ml_ops_api_client.models.framework import Framework as OpenapiFramework


class Framework(str, Enum):
    """
    An Enum to represent the supported frameworks.
    """

    # Models supported by upload / upload_stream endpoint
    ONNX = "onnx"
    TENSORFLOW = "tensorflow"
    PYTHON = "python"

    # Models supported by upload_and_convert endpoint
    KERAS = "keras"
    SKLEARN = "sklearn"
    PYTORCH = "pytorch"
    XGBOOST = "xgboost"
    HUGGING_FACE_FEATURE_EXTRACTION = "hugging-face-feature-extraction"
    HUGGING_FACE_IMAGE_CLASSIFICATION = "hugging-face-image-classification"
    HUGGING_FACE_IMAGE_SEGMENTATION = "hugging-face-image-segmentation"
    HUGGING_FACE_IMAGE_TO_TEXT = "hugging-face-image-to-text"
    HUGGING_FACE_OBJECT_DETECTION = "hugging-face-object-detection"
    HUGGING_FACE_QUESTION_ANSWERING = "hugging-face-question-answering"
    HUGGING_FACE_STABLE_DIFFUSION_TEXT_2_IMG = (
        "hugging-face-stable-diffusion-text-2-img"
    )
    HUGGING_FACE_SUMMARIZATION = "hugging-face-summarization"
    HUGGING_FACE_TEXT_CLASSIFICATION = "hugging-face-text-classification"
    HUGGING_FACE_TRANSLATION = "hugging-face-translation"
    HUGGING_FACE_ZERO_SHOT_CLASSIFICATION = "hugging-face-zero-shot-classification"
    HUGGING_FACE_ZERO_SHOT_IMAGE_CLASSIFICATION = (
        "hugging-face-zero-shot-image-classification"
    )
    HUGGING_FACE_ZERO_SHOT_OBJECT_DETECTION = "hugging-face-zero-shot-object-detection"
    HUGGING_FACE_SENTIMENT_ANALYSIS = "hugging-face-sentiment-analysis"
    HUGGING_FACE_TEXT_GENERATION = "hugging-face-text-generation"
    CUSTOM = "custom"

    def _to_openapi_framework(self) -> OpenapiFramework:
        """This is used to convert the enum into the autogenerated version.

        This is really for type-checking, since Python won't actually stop you,
        and the strings should be the same.
        """
        return OpenapiFramework[self.name]
