from enum import Enum
from typing import Dict, Optional, List
import turing.generated.models
from turing.generated.model_utils import OpenApiModel
from turing.generated.model.env_var import EnvVar
from .source import EnsemblingJobSource, EnsemblingJobPredictionSource
from .sink import EnsemblingJobSink
from ...mounted_mlp_secret import MountedMLPSecret

ResourceRequest = turing.generated.models.EnsemblingResources


class ResultType(Enum):
    DOUBLE = 0
    FLOAT = 1
    INTEGER = 2
    LONG = 3
    STRING = 4
    ARRAY = 10

    def to_open_api(self) -> OpenApiModel:
        return turing.generated.models.EnsemblingJobResultType(self.name)


class ResultConfig:
    def __init__(
        self, type: ResultType, column_name: str, item_type: Optional[ResultType] = None
    ):
        self._type, self._column_name, self._item_type = type, column_name, item_type

    def to_open_api(self) -> OpenApiModel:
        kwargs = {"type": self._type.to_open_api(), "column_name": self._column_name}
        if self._item_type:
            kwargs["item_type"] = self._item_type.to_open_api()

        return turing.generated.models.EnsemblingJobEnsemblerSpecResult(**kwargs)


class EnsemblingJobConfig:
    """
    Configuration of the batch ensembling job
    """

    def __init__(
        self,
        source: EnsemblingJobSource,
        predictions: Dict[str, EnsemblingJobPredictionSource],
        result_config: ResultConfig,
        sink: EnsemblingJobSink,
        service_account: str,
        resource_request: ResourceRequest = None,
        env_vars: Dict[str, str] = {},
        secrets: List[MountedMLPSecret] = None,
    ):
        """
        Create new instance of batch ensembling job configuration

        :param source: source configuration
        :param predictions: dictionary with configuration of model predictions
        :param result_config: configuration of ensembling results
        :param sink: sink configuration
        :param service_account:  secret name containing the service account for executing the ensembling job
        :param resource_request: optional resource request for starting the ensembling job.
            If not given the system default will be used.
        :param env_vars: optional environment variables in the form of a key value pair in a list.
        :param secrets: list of MLP secrets to mount into the ensembling job environment as environment variables
        """
        self._source = source
        self._predictions = predictions
        self._result_config = result_config
        self._sink = sink
        self._service_account = service_account
        self._resource_request = resource_request
        self._env_vars = env_vars
        self._secrets = secrets

    @property
    def source(self) -> "EnsemblingJobSource":
        return self._source

    @property
    def predictions(self) -> Dict[str, EnsemblingJobPredictionSource]:
        return self._predictions

    @property
    def sink(self) -> "EnsemblingJobSink":
        return self._sink

    @property
    def result_config(self) -> "ResultConfig":
        return self._result_config

    @property
    def service_account(self) -> str:
        return self._service_account

    @property
    def env_vars(self) -> Dict[str, str]:
        return self._env_vars

    @property
    def secrets(self) -> Optional[List[MountedMLPSecret]]:
        return self._secrets

    @property
    def resource_request(self) -> Optional["ResourceRequest"]:
        return self._resource_request

    def job_spec(self) -> turing.generated.models.EnsemblingJobSpec:
        return turing.generated.models.EnsemblingJobSpec(
            source=self.source.to_open_api(),
            predictions={
                name: source.to_open_api() for name, source in self.predictions.items()
            },
            ensembler=turing.generated.models.EnsemblingJobEnsemblerSpec(
                result=self.result_config.to_open_api()
            ),
            sink=self.sink.to_open_api(),
        )

    def infra_spec(self) -> turing.generated.models.EnsemblerInfraConfig:
        if self.env_vars is None:
            env_vars = []
        else:
            env_vars = [
                EnvVar(name=name, value=value) for name, value in self.env_vars.items()
            ]
        return turing.generated.models.EnsemblerInfraConfig(
            service_account_name=self.service_account,
            resources=self.resource_request,
            env=env_vars,
            secrets=(
                [secret.to_open_api() for secret in self.secrets]
                if self.secrets
                else []
            ),
        )
