"""
Copyright 2023 The HuggingFace Team
"""

from dataclasses import dataclass
from typing import List, Union

from autotrain import logger
from autotrain.backend import SpaceRunner
from autotrain.trainers.clm.params import LLMTrainingParams
from autotrain.trainers.dreambooth.params import DreamBoothTrainingParams
from autotrain.trainers.image_classification.params import ImageClassificationParams
from autotrain.trainers.seq2seq.params import Seq2SeqParams
from autotrain.trainers.tabular.params import TabularParams
from autotrain.trainers.text_classification.params import TextClassificationParams


@dataclass
class AutoTrainProject:
    params: Union[
        List[
            Union[
                LLMTrainingParams,
                TextClassificationParams,
                TabularParams,
                DreamBoothTrainingParams,
                Seq2SeqParams,
                ImageClassificationParams,
            ]
        ],
        LLMTrainingParams,
        TextClassificationParams,
        TabularParams,
        DreamBoothTrainingParams,
        Seq2SeqParams,
        ImageClassificationParams,
    ]
    backend: str

    def __post_init__(self):
        self.spaces_backends = {
            "A10G Large": "spaces-a10gl",
            "A10G Small": "spaces-a10gs",
            "A100 Large": "spaces-a100",
            "T4 Medium": "spaces-t4m",
            "T4 Small": "spaces-t4s",
            "CPU Upgrade": "spaces-cpu",
            "CPU (Free)": "spaces-cpuf",
            "DGX 1xA100": "dgx-a100",
            "DGX 2xA100": "dgx-2a100",
            "DGX 4xA100": "dgx-4a100",
            "DGX 8xA100": "dgx-8a100",
            "Local": "local",
            "EP US-East-1 1xA10g": "ep-aws-useast1-m",
            "EP US-East-1 1xA100": "ep-aws-useast1-xl",
            "EP US-East-1 2xA100": "ep-aws-useast1-2xl",
            "EP US-East-1 4xA100": "ep-aws-useast1-4xl",
            "EP US-East-1 8xA100": "ep-aws-useast1-8xl",
            "local": "local",
            "local-cli": "local-cli",
        }

    def create_spaces(self):
        sr = SpaceRunner(params=self.params, backend=self.spaces_backends[self.backend])
        space_id = sr.prepare()
        logger.info(f"Space created with id: {space_id}")
        return space_id

    def create(self):
        if self.backend == "AutoTrain":
            raise NotImplementedError
        if self.backend in self.spaces_backends:
            return self.create_spaces()
