# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Union

from oumi.core.configs import JobConfig
from oumi.core.launcher import BaseCloud, BaseCluster, JobStatus
from oumi.core.registry import REGISTRY, RegistryType



[docs]
class Launcher:
    """A class for managing the lifecycle of jobs on different clouds."""

    def __init__(self):
        """Initializes a new instance of the Launcher class."""
        self._clouds: dict[str, BaseCloud] = dict()
        self._initialize_new_clouds()

    def _initialize_new_clouds(self) -> None:
        """Initializes new clouds. Existing clouds are not re-initialized."""
        for name, builder in REGISTRY.get_all(RegistryType.CLOUD).items():
            if name not in self._clouds:
                self._clouds[name] = builder()

    def _get_cloud_by_name(self, cloud: str) -> BaseCloud:
        """Gets the cloud instance for the specified cloud name."""
        if cloud not in self._clouds:
            cloud_builder = REGISTRY.get(cloud, RegistryType.CLOUD)
            if not cloud_builder:
                raise ValueError(f"Cloud {cloud} not found in the registry.")
            self._clouds[cloud] = cloud_builder()
        return self._clouds[cloud]


[docs]
    def cancel(self, job_id: str, cloud_name: str, cluster_name: str) -> JobStatus:
        """Cancels the specified job."""
        cloud = self._get_cloud_by_name(cloud_name)
        cluster = cloud.get_cluster(cluster_name)
        if not cluster:
            raise ValueError(f"Cluster {cluster_name} not found.")
        return cluster.cancel_job(job_id)



[docs]
    def down(self, cloud_name: str, cluster_name: str) -> None:
        """Turns down the specified cluster."""
        cloud = self._get_cloud_by_name(cloud_name)
        cluster = cloud.get_cluster(cluster_name)
        if not cluster:
            raise ValueError(f"Cluster {cluster_name} not found.")
        cluster.down()



[docs]
    def get_cloud(self, job_or_cloud: Union[JobConfig, str]) -> BaseCloud:
        """Gets the cloud instance for the specified job."""
        if isinstance(job_or_cloud, str):
            return self._get_cloud_by_name(job_or_cloud)
        return self._get_cloud_by_name(job_or_cloud.resources.cloud)



[docs]
    def run(self, job: JobConfig, cluster_name: str) -> JobStatus:
        """Runs the specified job on the specified cluster.

        Args:
            job: The job configuration.
            cluster_name: The name of the cluster to run the job on.

        Returns:
            Optional[JobStatus]: The status of the job.
        """
        cloud = self.get_cloud(job)
        cluster = cloud.get_cluster(cluster_name)
        if not cluster:
            raise ValueError(f"Cluster {cluster_name} not found.")
        return cluster.run_job(job)



[docs]
    def status(
        self,
        cloud: Optional[str] = None,
        cluster: Optional[str] = None,
        id: Optional[str] = None,
    ) -> dict[str, list[JobStatus]]:
        """Gets the status of all jobs across all clusters.

        Args:
            cloud: If specified, filters all jobs to only those on the specified cloud.
            cluster: If specified, filters all jobs to only those on the specified
                cluster.
            id: If specified, filters all jobs to only those with the specified ID.

        Returns:
            dict[str, list(JobStatus)]: The status of all jobs, indexed by cloud name.
        """
        # Pick up any newly registered cloud builders.
        self._initialize_new_clouds()
        statuses: dict[str, list[JobStatus]] = {}
        for cloud_name, target_cloud in self._clouds.items():
            # Ignore clouds not matching the filter criteria.
            if cloud and cloud_name != cloud:
                continue
            statuses[cloud_name] = []
            for target_cluster in target_cloud.list_clusters():
                # Ignore clusters not matching the filter criteria.
                if cluster and target_cluster.name() != cluster:
                    continue
                for job in target_cluster.get_jobs():
                    # Ignore jobs not matching the filter criteria.
                    if id and job.id != id:
                        continue
                    statuses[cloud_name].append(job)
        return statuses



[docs]
    def stop(self, cloud_name: str, cluster_name: str) -> None:
        """Stops the specified cluster."""
        cloud = self._get_cloud_by_name(cloud_name)
        cluster = cloud.get_cluster(cluster_name)
        if not cluster:
            raise ValueError(f"Cluster {cluster_name} not found.")
        cluster.stop()



[docs]
    def up(
        self, job: JobConfig, cluster_name: Optional[str], **kwargs
    ) -> tuple[BaseCluster, JobStatus]:
        """Creates a new cluster and starts the specified job on it."""
        cloud = self.get_cloud(job)
        job_status = cloud.up_cluster(job, cluster_name, **kwargs)
        cluster = cloud.get_cluster(job_status.cluster)
        if not cluster:
            raise RuntimeError(f"Cluster {job_status.cluster} not found.")
        return (cluster, job_status)



[docs]
    def which_clouds(self) -> list[str]:
        """Gets the names of all clouds in the registry."""
        return [name for name, _ in REGISTRY.get_all(RegistryType.CLOUD).items()]




LAUNCHER = Launcher()
# Explicitly expose the public methods of the default Launcher instance.
#: A convenience method for Launcher.cancel.
cancel = LAUNCHER.cancel
#: A convenience method for Launcher.down.
down = LAUNCHER.down
#: A convenience method for Launcher.get_cloud.
get_cloud = LAUNCHER.get_cloud
#: A convenience method for Launcher.run.
run = LAUNCHER.run
#: A convenience method for Launcher.status.
status = LAUNCHER.status
#: A convenience method for Launcher.stop.
stop = LAUNCHER.stop
#: A convenience method for Launcher.up.
up = LAUNCHER.up
#: A convenience method for Launcher.which_clouds.
which_clouds = LAUNCHER.which_clouds