#  Copyright (c) ZenML GmbH 2022. 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.
"""Utility functions for building manifests for k8s pods."""

import base64
import os
import sys
from typing import Any, Dict, List, Mapping, Optional

from kubernetes import client as k8s_client

from zenml.client import Client
from zenml.config.global_config import GlobalConfiguration
from zenml.constants import ENV_ZENML_ENABLE_REPO_INIT_WARNINGS
from zenml.integrations.airflow.orchestrators.dag_generator import (
    ENV_ZENML_LOCAL_STORES_PATH,
)
from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
from zenml.logger import get_logger
from zenml.utils import string_utils

logger = get_logger(__name__)


def add_local_stores_mount(
    pod_spec: k8s_client.V1PodSpec,
) -> None:
    """Makes changes in place to the configuration of the pod spec.

    Configures mounted volumes for stack components that write to a local
    path.

    Args:
        pod_spec: The pod spec to update.
    """
    assert len(pod_spec.containers) == 1
    container_spec: k8s_client.V1Container = pod_spec.containers[0]

    stack = Client().active_stack

    stack.check_local_paths()

    local_stores_path = GlobalConfiguration().local_stores_path

    host_path = k8s_client.V1HostPathVolumeSource(
        path=local_stores_path, type="Directory"
    )

    pod_spec.volumes = pod_spec.volumes or []
    pod_spec.volumes.append(
        k8s_client.V1Volume(
            name="local-stores",
            host_path=host_path,
        )
    )
    container_spec.volume_mounts = container_spec.volume_mounts or []
    container_spec.volume_mounts.append(
        k8s_client.V1VolumeMount(
            name="local-stores",
            mount_path=local_stores_path,
        )
    )

    if sys.platform == "win32":
        # File permissions are not checked on Windows. This if clause
        # prevents mypy from complaining about unused 'type: ignore'
        # statements
        pass
    else:
        # Run KFP containers in the context of the local UID/GID
        # to ensure that the local stores can be shared
        # with the local pipeline runs.
        pod_spec.security_context = k8s_client.V1SecurityContext(
            run_as_user=os.getuid(),
            run_as_group=os.getgid(),
        )

    container_spec.env = container_spec.env or []
    container_spec.env.append(
        k8s_client.V1EnvVar(
            name=ENV_ZENML_LOCAL_STORES_PATH,
            value=local_stores_path,
        )
    )


def build_pod_manifest(
    pod_name: Optional[str],
    image_name: str,
    command: List[str],
    args: List[str],
    privileged: bool,
    pod_settings: Optional[KubernetesPodSettings] = None,
    service_account_name: Optional[str] = None,
    env: Optional[Dict[str, str]] = None,
    labels: Optional[Dict[str, str]] = None,
    mount_local_stores: bool = False,
    owner_references: Optional[List[k8s_client.V1OwnerReference]] = None,
    termination_grace_period_seconds: Optional[int] = 30,
) -> k8s_client.V1Pod:
    """Build a Kubernetes pod manifest for a ZenML run or step.

    Args:
        pod_name: Name of the pod.
        image_name: Name of the Docker image.
        command: Command to execute the entrypoint in the pod.
        args: Arguments provided to the entrypoint command.
        privileged: Whether to run the container in privileged mode.
        pod_settings: Optional settings for the pod.
        service_account_name: Optional name of a service account.
            Can be used to assign certain roles to a pod, e.g., to allow it to
            run Kubernetes commands from within the cluster.
        env: Environment variables to set.
        labels: Labels to add to the pod.
        mount_local_stores: Whether to mount the local stores path inside the
            pod.
        owner_references: List of owner references for the pod.
        termination_grace_period_seconds: The amount of seconds to wait for a
            pod to shutdown gracefully.

    Returns:
        Pod manifest.
    """
    env = env.copy() if env else {}
    env.setdefault(ENV_ZENML_ENABLE_REPO_INIT_WARNINGS, "False")

    security_context = k8s_client.V1SecurityContext(privileged=privileged)
    container_spec = k8s_client.V1Container(
        name="main",
        image=image_name,
        command=command,
        args=args,
        env=[
            k8s_client.V1EnvVar(name=name, value=value)
            for name, value in env.items()
        ],
        security_context=security_context,
    )
    image_pull_secrets = []
    if pod_settings:
        image_pull_secrets = [
            k8s_client.V1LocalObjectReference(name=name)
            for name in pod_settings.image_pull_secrets
        ]

    pod_spec = k8s_client.V1PodSpec(
        containers=[container_spec],
        restart_policy="Never",
        image_pull_secrets=image_pull_secrets,
        termination_grace_period_seconds=termination_grace_period_seconds,
    )

    if service_account_name is not None:
        pod_spec.service_account_name = service_account_name

    # Apply pod settings if provided
    labels = labels or {}

    if pod_settings:
        substitutions = {
            "{{ image }}": image_name,
        }
        add_pod_settings(
            pod_spec=pod_spec,
            settings=pod_settings,
            substitutions=substitutions,
        )

    if pod_settings and pod_settings.labels:
        labels.update(pod_settings.labels)

    pod_metadata = k8s_client.V1ObjectMeta(
        name=pod_name,
        labels=labels,
        owner_references=owner_references,
    )

    if pod_settings and pod_settings.annotations:
        pod_metadata.annotations = pod_settings.annotations

    pod_manifest = k8s_client.V1Pod(
        kind="Pod",
        api_version="v1",
        metadata=pod_metadata,
        spec=pod_spec,
    )

    if mount_local_stores:
        add_local_stores_mount(pod_spec)

    return pod_manifest


def add_pod_settings(
    pod_spec: k8s_client.V1PodSpec,
    settings: KubernetesPodSettings,
    substitutions: Dict[str, str],
) -> None:
    """Updates pod `spec` fields in place if passed in orchestrator settings.

    Args:
        pod_spec: Pod spec to update.
        settings: Pod settings to apply.
        substitutions: Substitutions to apply to additional pod spec arguments.
    """
    if settings.node_selectors:
        pod_spec.node_selector = settings.node_selectors

    if settings.affinity:
        pod_spec.affinity = settings.affinity

    if settings.tolerations:
        pod_spec.tolerations = settings.tolerations

    for container in pod_spec.containers:
        assert isinstance(container, k8s_client.V1Container)
        container._resources = settings.resources
        if settings.volume_mounts:
            if container.volume_mounts:
                container.volume_mounts.extend(settings.volume_mounts)
            else:
                container.volume_mounts = settings.volume_mounts

        if settings.env:
            if container.env:
                container.env.extend(settings.env)
            else:
                container.env = settings.env

        if settings.env_from:
            if container.env_from:
                container.env_from.extend(settings.env_from)
            else:
                container.env_from = settings.env_from

        if settings.container_security_context:
            if container.security_context:
                existing_dict = container.security_context.to_dict()
                existing_dict.update(settings.container_security_context)
                container.security_context = existing_dict
            else:
                container.security_context = (
                    settings.container_security_context
                )

    if settings.volumes:
        if pod_spec.volumes:
            pod_spec.volumes.extend(settings.volumes)
        else:
            pod_spec.volumes = settings.volumes

    if settings.host_ipc:
        pod_spec.host_ipc = settings.host_ipc

    if settings.scheduler_name:
        pod_spec.scheduler_name = settings.scheduler_name

    for key, value in settings.additional_pod_spec_args.items():
        if not hasattr(pod_spec, key):
            logger.warning(f"Ignoring invalid Pod Spec argument `{key}`.")
        else:
            if value is None:
                continue

            def _substitute(v: str) -> str:
                return substitutions.get(v, v)

            value = string_utils.substitute_string(value, _substitute)

            existing_value = getattr(pod_spec, key)
            if isinstance(existing_value, list):
                existing_value.extend(value)
            elif isinstance(existing_value, dict):
                existing_value.update(value)
            else:
                setattr(pod_spec, key, value)


def build_role_binding_manifest_for_service_account(
    name: str,
    role_name: str,
    service_account_name: str,
    namespace: str = "default",
) -> Dict[str, Any]:
    """Build a manifest for a role binding of a service account.

    Args:
        name: Name of the cluster role binding.
        role_name: Name of the role.
        service_account_name: Name of the service account.
        namespace: Kubernetes namespace. Defaults to "default".

    Returns:
        Manifest for a cluster role binding of a service account.
    """
    return {
        "apiVersion": "rbac.authorization.k8s.io/v1",
        "kind": "RoleBinding",
        "metadata": {"name": name},
        "subjects": [
            {
                "kind": "ServiceAccount",
                "name": service_account_name,
                "namespace": namespace,
            }
        ],
        "roleRef": {
            "kind": "ClusterRole",
            "name": role_name,
            "apiGroup": "rbac.authorization.k8s.io",
        },
    }


def build_service_account_manifest(
    name: str, namespace: str = "default"
) -> Dict[str, Any]:
    """Build the manifest for a service account.

    Args:
        name: Name of the service account.
        namespace: Kubernetes namespace. Defaults to "default".

    Returns:
        Manifest for a service account.
    """
    return {
        "apiVersion": "v1",
        "metadata": {
            "name": name,
            "namespace": namespace,
        },
    }


def build_namespace_manifest(
    namespace: str, labels: Optional[Dict[str, str]] = None
) -> Dict[str, Any]:
    """Build the manifest for a new namespace.

    Args:
        namespace: Kubernetes namespace.
        labels: Optional labels to apply to the namespace. These labels are used
            to track namespace ownership for cleanup. For example, adding
            'zenml-deployment-id' allows us to identify and safely delete
            namespaces created by ZenML deployments.

    Returns:
        Manifest of the new namespace.
    """
    metadata: Dict[str, Any] = {"name": namespace}
    if labels:
        metadata["labels"] = labels

    return {
        "apiVersion": "v1",
        "kind": "Namespace",
        "metadata": metadata,
    }


def build_secret_manifest(
    name: str,
    data: Mapping[str, Optional[str]],
    namespace: Optional[str] = None,
    secret_type: str = "Opaque",
) -> Dict[str, Any]:
    """Builds a Kubernetes secret manifest.

    Args:
        name: Name of the secret.
        data: The secret data.
        namespace: The namespace to create the secret in.
        secret_type: The secret type.

    Returns:
        The secret manifest.
    """
    encoded_data = {
        key: base64.b64encode(value.encode()).decode() if value else None
        for key, value in data.items()
    }

    return {
        "apiVersion": "v1",
        "kind": "Secret",
        "metadata": {"name": name, "namespace": namespace},
        "type": secret_type,
        "data": encoded_data,
    }


def pod_template_manifest_from_pod(
    pod: k8s_client.V1Pod,
) -> k8s_client.V1PodTemplateSpec:
    """Build a Kubernetes pod template manifest from a pod.

    Args:
        pod: The pod manifest to build the template from.

    Returns:
        The pod template manifest.
    """
    return k8s_client.V1PodTemplateSpec(
        metadata=pod.metadata,
        spec=pod.spec,
    )


def build_job_manifest(
    job_name: str,
    pod_template: k8s_client.V1PodTemplateSpec,
    backoff_limit: Optional[int] = None,
    ttl_seconds_after_finished: Optional[int] = None,
    labels: Optional[Dict[str, str]] = None,
    annotations: Optional[Dict[str, str]] = None,
    active_deadline_seconds: Optional[int] = None,
    pod_failure_policy: Optional[Dict[str, Any]] = None,
    owner_references: Optional[List[k8s_client.V1OwnerReference]] = None,
) -> k8s_client.V1Job:
    """Build a Kubernetes job manifest.

    Args:
        job_name: Name of the job.
        pod_template: The pod template to use for the job.
        backoff_limit: The backoff limit for the job.
        ttl_seconds_after_finished: The TTL seconds after finished for the job.
        labels: The labels to use for the job.
        annotations: The annotations to use for the job.
        active_deadline_seconds: The active deadline seconds for the job.
        pod_failure_policy: The pod failure policy for the job.
        owner_references: The owner references for the job.

    Returns:
        The Kubernetes job manifest.
    """
    job_spec = k8s_client.V1JobSpec(
        template=pod_template,
        backoff_limit=backoff_limit,
        parallelism=1,
        ttl_seconds_after_finished=ttl_seconds_after_finished,
        active_deadline_seconds=active_deadline_seconds,
        pod_failure_policy=pod_failure_policy,
    )

    job_metadata = k8s_client.V1ObjectMeta(
        name=job_name,
        labels=labels,
        annotations=annotations,
        owner_references=owner_references,
    )

    return k8s_client.V1Job(spec=job_spec, metadata=job_metadata)


def job_template_manifest_from_job(
    job: k8s_client.V1Job,
) -> k8s_client.V1JobTemplateSpec:
    """Build a Kubernetes job template manifest from a job.

    Args:
        job: The job manifest to build the template from.

    Returns:
        The job template manifest.
    """
    return k8s_client.V1JobTemplateSpec(
        metadata=job.metadata,
        spec=job.spec,
    )


def build_cron_job_manifest(
    job_template: k8s_client.V1JobTemplateSpec,
    cron_expression: str,
    successful_jobs_history_limit: Optional[int] = None,
    failed_jobs_history_limit: Optional[int] = None,
    concurrency_policy: Optional[str] = None,
    starting_deadline_seconds: Optional[int] = None,
) -> k8s_client.V1CronJob:
    """Build a Kubernetes cron job manifest.

    Args:
        job_template: The job template to use for the cron job.
        cron_expression: The cron expression to use for the cron job.
        successful_jobs_history_limit: The number of successful jobs to keep.
        failed_jobs_history_limit: The number of failed jobs to keep.
        concurrency_policy: The concurrency policy for the cron job
            (Allow, Forbid, or Replace).
        starting_deadline_seconds: The deadline in seconds for starting
            a job if it misses its scheduled time.

    Returns:
        The Kubernetes cron job manifest.
    """
    spec = k8s_client.V1CronJobSpec(
        schedule=cron_expression,
        successful_jobs_history_limit=successful_jobs_history_limit,
        failed_jobs_history_limit=failed_jobs_history_limit,
        concurrency_policy=concurrency_policy,
        starting_deadline_seconds=starting_deadline_seconds,
        job_template=job_template,
    )

    return k8s_client.V1CronJob(
        kind="CronJob",
        api_version="batch/v1",
        metadata=job_template.metadata,
        spec=spec,
    )
