from harbor.metrics.base import BaseMetric
from harbor.models.job.config import RetryConfig
from harbor.models.orchestrator_type import OrchestratorType
from harbor.models.trial.config import TrialConfig
from harbor.orchestrators.base import BaseOrchestrator
from harbor.orchestrators.local import LocalOrchestrator


class OrchestratorFactory:
    _ORCHESTRATORS: list[type[BaseOrchestrator]] = [LocalOrchestrator]
    _ORCHESTRATOR_MAP: dict[OrchestratorType, type[BaseOrchestrator]] = {
        orchestrator.type(): orchestrator for orchestrator in _ORCHESTRATORS
    }

    @classmethod
    def create_orchestrator(
        cls,
        orchestrator_type: OrchestratorType,
        trial_configs: list[TrialConfig],
        n_concurrent_trials: int,
        metrics: dict[str, list[BaseMetric]],
        quiet: bool = False,
        retry_config: RetryConfig | None = None,
        **kwargs,
    ) -> BaseOrchestrator:
        if orchestrator_type not in cls._ORCHESTRATOR_MAP:
            raise ValueError(
                f"Unsupported orchestrator type: {orchestrator_type}. This could be "
                "because the orchestrator is not registered in the OrchestratorFactory "
                "or because the orchestrator type is invalid."
            )

        orchestrator_class = cls._ORCHESTRATOR_MAP[orchestrator_type]

        return orchestrator_class(
            trial_configs, n_concurrent_trials, metrics, quiet, retry_config, **kwargs
        )
