#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
#
#  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:
#
#       https://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.
"""Skypilot orchestrator base config and settings."""

from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import Field, model_validator

from zenml.config.base_settings import BaseSettings
from zenml.logger import get_logger
from zenml.orchestrators import BaseOrchestratorConfig

logger = get_logger(__name__)


class SkypilotBaseOrchestratorSettings(BaseSettings):
    """Skypilot orchestrator base settings.

    Attributes:
        instance_type: the instance type to use.
        cpus: the number of CPUs required for the task.
            If a str, must be a string of the form `'2'` or `'2+'`, where
            the `+` indicates that the task requires at least 2 CPUs.
        memory: the amount of memory in GiB required. If a
            str, must be a string of the form `'16'` or `'16+'`, where
            the `+` indicates that the task requires at least 16 GB of memory.
        accelerators: the accelerators required. If a str, must be
            a string of the form `'V100'` or `'V100:2'`, where the `:2`
            indicates that the task requires 2 V100 GPUs. If a dict, must be a
            dict of the form `{'V100': 2}` or `{'tpu-v2-8': 1}`.
        accelerator_args: accelerator-specific arguments. For example,
            `{'tpu_vm': True, 'runtime_version': 'tpu-vm-base'}` for TPUs.
        use_spot: whether to use spot instances. If None, defaults to
            False.
        job_recovery: the spot recovery strategy to use for the managed
            spot to recover the cluster from preemption. Refer to
            `recovery_strategy module <https://github.com/skypilot-org/skypilot/blob/master/sky/spot/recovery_strategy.py>`__ # pylint: disable=line-too-long
            for more details.
        region: the region to use.
        zone: the zone to use.
        image_id: the image ID to use. If a str, must be a string
            of the image id from the cloud, such as AWS:
            ``'ami-1234567890abcdef0'``, GCP:
            ``'projects/my-project-id/global/images/my-image-name'``;
            Or, a image tag provided by SkyPilot, such as AWS:
            ``'skypilot:gpu-ubuntu-2004'``. If a dict, must be a dict mapping
            from region to image ID, such as:

            .. code-block:: python

                {
                'us-west1': 'ami-1234567890abcdef0',
                'us-east1': 'ami-1234567890abcdef0'
                }

        disk_size: the size of the OS disk in GiB.
        disk_tier: the disk performance tier to use. If None, defaults to
            ``'medium'``.
        network_tier: the network performance tier to use where supported by
            the cloud provider (e.g., ``'standard'`` or ``'best'``).
        infra: SkyPilot infrastructure selector (e.g., ``"aws/us-east-1"``,
            ``"gcp/us-central1-a"``, ``"k8s/my-cluster-ctx"``). Mutually
            exclusive with explicit cloud/region/zone selection.
        num_nodes: number of nodes for multi-node jobs.
        ports: Ports to expose. Could be an integer, a range, or a list of
            integers and ranges. All ports will be exposed to the public internet.
        labels: Labels to apply to instances as key-value pairs. These are
            mapped to cloud-specific implementations (instance tags in AWS,
            instance labels in GCP, etc.)
        any_of: List of candidate resources to try in order of preference based on
            cost (determined by the optimizer).
        ordered: List of candidate resources to try in the specified order.

        cluster_name: name of the cluster to create/reuse.  If None,
            auto-generate a name.
        retry_until_up: whether to retry launching the cluster until it is
            up.
        idle_minutes_to_autostop: automatically stop the cluster after this
            many minute of idleness, i.e., no running or pending jobs in the
            cluster's job queue. Idleness gets reset whenever setting-up/
            running/pending jobs are found in the job queue. Setting this
            flag schedules an autostop after the launch completes. If not set,
            the cluster will not be autostopped.
        down: Tear down the cluster after all jobs finish (successfully or
            abnormally). If --idle-minutes-to-autostop is also set, the
            cluster will be torn down after the specified idle time.
            Note that if errors occur during provisioning/data syncing/setting
            up, the cluster will not be torn down for debugging purposes.
        stream_logs: if True, show the logs in the terminal.
        docker_run_args: Optional arguments to pass to the `docker run` command
            running inside the VM.
        workdir: Working directory to sync to the VM. Synced to ~/sky_workdir.
        task_name: Task name used for display purposes.
        file_mounts: File and storage mounts configuration for remote cluster.
        envs: Environment variables for the task.
        task_settings: Dictionary of arbitrary settings to pass to sky.Task().
            This allows passing future parameters added by SkyPilot without
            requiring updates to ZenML.
        resources_settings: Dictionary of arbitrary settings to pass to
            sky.Resources(). This allows passing future parameters added
            by SkyPilot without requiring updates to ZenML.
        launch_settings: Dictionary of arbitrary settings to pass to
            sky.launch(). This allows passing future parameters added
            by SkyPilot without requiring updates to ZenML.
    """

    # Resources
    instance_type: Optional[str] = None
    cpus: Union[None, int, float, str] = Field(
        default=None, union_mode="left_to_right"
    )
    memory: Union[None, int, float, str] = Field(
        default=None, union_mode="left_to_right"
    )
    accelerators: Union[None, str, Dict[str, int]] = Field(
        default=None, union_mode="left_to_right"
    )
    accelerator_args: Optional[Dict[str, str]] = None
    use_spot: Optional[bool] = None
    job_recovery: Union[None, str, Dict[str, Any]] = Field(
        default=None, union_mode="left_to_right"
    )
    region: Optional[str] = None
    zone: Optional[str] = None
    image_id: Union[Dict[str, str], str, None] = Field(
        default=None, union_mode="left_to_right"
    )
    disk_size: Optional[int] = None
    disk_tier: Optional[Literal["high", "medium", "low", "ultra", "best"]] = (
        None
    )
    network_tier: Optional[str] = None
    infra: Optional[str] = None
    num_nodes: Optional[int] = None

    # Run settings
    cluster_name: Optional[str] = None
    retry_until_up: bool = False
    idle_minutes_to_autostop: Optional[int] = 30
    down: bool = True
    stream_logs: bool = True
    docker_run_args: List[str] = Field(default_factory=list)

    # Additional SkyPilot features
    ports: Union[None, int, str, List[Union[int, str]]] = Field(
        default=None, union_mode="left_to_right"
    )
    labels: Optional[Dict[str, str]] = None
    any_of: Optional[List[Dict[str, Any]]] = None
    ordered: Optional[List[Dict[str, Any]]] = None
    workdir: Optional[str] = None
    task_name: Optional[str] = None
    file_mounts: Optional[Dict[str, Any]] = None
    envs: Optional[Dict[str, str]] = None

    # Future-proofing settings dictionaries
    task_settings: Dict[str, Any] = Field(default_factory=dict)
    resources_settings: Dict[str, Any] = Field(default_factory=dict)
    launch_settings: Dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="after")
    def _validate_infra_configuration(
        self,
    ) -> "SkypilotBaseOrchestratorSettings":
        """Validates settings that are incompatible with infra shortcuts.

        Returns:
            The validated settings object.

        Raises:
            ValueError: If unsupported combinations with `infra` are set.
        """
        if self.infra is None:
            return self

        conflicts = []
        if self.region is not None:
            conflicts.append("region")
        if self.zone is not None:
            conflicts.append("zone")
        for key in ("cloud", "region", "zone"):
            if key in self.resources_settings:
                conflicts.append(f"resources_settings['{key}']")

        if conflicts:
            raise ValueError(
                "The `infra` setting cannot be combined with explicit cloud/"
                "region/zone settings. Remove the following conflicting "
                f"values: {', '.join(conflicts)}."
            )

        return self


class SkypilotBaseOrchestratorConfig(
    BaseOrchestratorConfig, SkypilotBaseOrchestratorSettings
):
    """Skypilot orchestrator base config.

    Attributes:
        disable_step_based_settings: whether to disable step-based settings.
            If True, the orchestrator will run all steps with the pipeline
            settings in one single VM. If False, the orchestrator will run
            each step with its own settings in separate VMs if provided.
    """

    disable_step_based_settings: bool = False

    @property
    def is_local(self) -> bool:
        """Checks if this stack component is running locally.

        Returns:
            True if this config is for a local component, False otherwise.
        """
        return False

    @property
    def supports_client_side_caching(self) -> bool:
        """Whether the orchestrator supports client side caching.

        Returns:
            Whether the orchestrator supports client side caching.
        """
        # The Skypilot orchestrator runs the entire pipeline in a single VM, or
        # starts additional VMs from the root VM. Both of those cases are
        # currently not supported when using client-side caching.
        return False
