import argparse
import json
import logging
import os
import re
import shlex
import socket
import sys
import uuid
from typing import Any
from typing import Dict
from typing import List
from typing import Mapping
from typing import Optional
from typing import Tuple
from typing import Union

import yaml
from service_configuration_lib import read_service_configuration
from service_configuration_lib import spark_config
from service_configuration_lib.spark_config import get_aws_credentials
from service_configuration_lib.spark_config import get_grafana_url
from service_configuration_lib.spark_config import get_resources_requested
from service_configuration_lib.spark_config import get_signalfx_url
from service_configuration_lib.spark_config import get_spark_hourly_cost
from service_configuration_lib.spark_config import UnsupportedClusterManagerException

from paasta_tools.cli.cmds.check import makefile_responds_to
from paasta_tools.cli.cmds.cook_image import paasta_cook_image
from paasta_tools.cli.utils import get_instance_config
from paasta_tools.cli.utils import lazy_choices_completer
from paasta_tools.cli.utils import list_instances
from paasta_tools.clusterman import get_clusterman_metrics
from paasta_tools.kubernetes_tools import limit_size_with_hash
from paasta_tools.spark_tools import DEFAULT_SPARK_SERVICE
from paasta_tools.spark_tools import get_volumes_from_spark_k8s_configs
from paasta_tools.spark_tools import get_webui_url
from paasta_tools.spark_tools import inject_spark_conf_str
from paasta_tools.utils import _run
from paasta_tools.utils import DEFAULT_SOA_DIR
from paasta_tools.utils import get_possible_launched_by_user_variable_from_env
from paasta_tools.utils import get_username
from paasta_tools.utils import InstanceConfig
from paasta_tools.utils import list_services
from paasta_tools.utils import load_system_paasta_config
from paasta_tools.utils import NoConfigurationForServiceError
from paasta_tools.utils import NoDeploymentsAvailable
from paasta_tools.utils import NoDockerImageError
from paasta_tools.utils import PaastaColors
from paasta_tools.utils import PaastaNotConfiguredError
from paasta_tools.utils import SystemPaastaConfig


DEFAULT_AWS_REGION = "us-west-2"
DEFAULT_SPARK_WORK_DIR = "/spark_driver"
DEFAULT_SPARK_DOCKER_IMAGE_PREFIX = "paasta-spark-run"
DEFAULT_SPARK_DOCKER_REGISTRY = "docker-dev.yelpcorp.com"
SENSITIVE_ENV = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]
clusterman_metrics, CLUSTERMAN_YAML_FILE_PATH = get_clusterman_metrics()
CLUSTER_MANAGER_K8S = "kubernetes"
CLUSTER_MANAGER_LOCAL = "local"
CLUSTER_MANAGERS = {CLUSTER_MANAGER_K8S, CLUSTER_MANAGER_LOCAL}
DEFAULT_DOCKER_SHM_SIZE = "64m"
# Reference: https://spark.apache.org/docs/latest/configuration.html#application-properties
DEFAULT_DRIVER_CORES_BY_SPARK = 1
DEFAULT_DRIVER_MEMORY_BY_SPARK = "1g"
# Extra room for memory overhead and for any other running inside container
DOCKER_RESOURCE_ADJUSTMENT_FACTOR = 2

POD_TEMPLATE_DIR = "/nail/tmp"
POD_TEMPLATE_PATH = "/nail/tmp/spark-pt-{file_uuid}.yaml"
DEFAULT_RUNTIME_TIMEOUT = "12h"

POD_TEMPLATE = """
apiVersion: v1
kind: Pod
metadata:
  labels:
    spark: {spark_pod_label}
spec:
  dnsPolicy: Default
  affinity:
    podAffinity:
      preferredDuringSchedulingIgnoredDuringExecution:
      - weight: 95
        podAffinityTerm:
          labelSelector:
            matchExpressions:
            - key: spark
              operator: In
              values:
              - {spark_pod_label}
          topologyKey: topology.kubernetes.io/hostname
"""

deprecated_opts = {
    "j": "spark.jars",
    "jars": "spark.jars",
    "max-cores": "spark.cores.max",
    "executor-cores": "spark.executor.cores",
    "executor-memory": "spark.executor.memory",
    "driver-max-result-size": "spark.driver.maxResultSize",
    "driver-cores": "spark.driver.cores",
    "driver-memory": "spark.driver.memory",
}

SPARK_COMMANDS = {"pyspark", "spark-submit"}

log = logging.getLogger(__name__)


class DeprecatedAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        print(
            PaastaColors.red(
                "Use of {} is deprecated. Please use {}=value in --spark-args.".format(
                    option_string, deprecated_opts[option_string.strip("-")]
                )
            )
        )
        sys.exit(1)


def add_subparser(subparsers):
    list_parser = subparsers.add_parser(
        "spark-run",
        help="Run Spark on the PaaSTA cluster",
        description=(
            "'paasta spark-run' launches a Spark cluster on PaaSTA. "
            "It analyzes soa-configs and command line arguments to invoke "
            "a 'docker run'. By default, it will pull the Spark service "
            "image from the registry unless the --build option is used.\n\n"
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    group = list_parser.add_mutually_exclusive_group()
    group.add_argument(
        "-b",
        "--build",
        help="Build the docker image from scratch using the local Makefile's cook-image target.",
        action="store_true",
        default=False,
    )
    group.add_argument(
        "-I",
        "--image",
        help="Use the provided image to start the Spark driver and executors.",
    )

    list_parser.add_argument(
        "-e",
        "--enable-compact-bin-packing",
        help=(
            "Enabling compact bin packing will try to ensure executors are scheduled on the same nodes. Requires --cluster-manager to be kubernetes."
            " Always true by default, keep around for backward compability."
        ),
        action="store_true",
        default=True,
    )
    list_parser.add_argument(
        "--disable-compact-bin-packing",
        help=(
            "Disable compact bin packing. Requires --cluster-manager to be kubernetes. Note: this option is only for advanced Spark configurations,"
            " don't use it unless you've been instructed to do so."
        ),
        action="store_true",
        default=False,
    )
    list_parser.add_argument(
        "--docker-memory-limit",
        help=(
            "Set docker memory limit. Should be greater than driver memory. Defaults to 2x spark.driver.memory. Example: 2g, 500m, Max: 64g"
            " Note: If memory limit provided is greater than associated with the batch instance, it will default to max memory of the box."
        ),
        default=None,
    )
    list_parser.add_argument(
        "--docker-cpu-limit",
        help=(
            "Set docker cpus limit. Should be greater than driver cores. Defaults to 1x spark.driver.cores."
            " Note: The job will fail if the limit provided is greater than number of cores present on batch box (8 for production batch boxes)."
        ),
        default=None,
    )

    list_parser.add_argument(
        "--docker-shm-size",
        help=(
            "Set docker shared memory size limit for the driver's container. This is the same as setting docker run --shm-size and the shared"
            " memory is mounted to /dev/shm in the container. Anything written to the shared memory mount point counts towards the docker memory"
            " limit for the driver's container. Therefore, this should be less than --docker-memory-limit."
            f" Defaults to {DEFAULT_DOCKER_SHM_SIZE}. Example: 8g, 256m"
            " Note: this option is mainly useful when training TensorFlow models in the driver, with multiple GPUs using NCCL. The shared memory"
            f" space is used to sync gradient updates between GPUs during training. The default value of {DEFAULT_DOCKER_SHM_SIZE} is typically not large enough for"
            " this inter-gpu communication to run efficiently. We recommend a starting value of 8g to ensure that the entire set of model parameters"
            " can fit in the shared memory. This can be less if you are training a smaller model (<1g parameters) or more if you are using a larger model (>2.5g parameters)"
            " If you are observing low, average GPU utilization during epoch training (<65-70 percent) you can also try increasing this value; you may be"
            " resource constrained when GPUs sync training weights between mini-batches (there are other potential bottlenecks that could cause this as well)."
            " A tool such as nvidia-smi can be use to check GPU utilization."
            " This option also adds the --ulimit memlock=-1 to the docker run command since this is recommended for TensorFlow applications that use NCCL."
            " Please refer to docker run documentation for more details on --shm-size and --ulimit memlock=-1."
        ),
        default=None,
    )
    list_parser.add_argument(
        "--force-spark-resource-configs",
        help=(
            "Skip the resource/instances recalculation. "
            "This is strongly not recommended."
        ),
        action="store_true",
        default=False,
    )
    list_parser.add_argument(
        "--docker-registry",
        help="Docker registry to push the Spark image built.",
        default=None,
    )

    list_parser.add_argument(
        "-s",
        "--service",
        help="The name of the service from which the Spark image is built.",
        default=DEFAULT_SPARK_SERVICE,
    ).completer = lazy_choices_completer(list_services)

    list_parser.add_argument(
        "-i",
        "--instance",
        help="Start a docker run for a particular instance of the service.",
        default="adhoc",
    ).completer = lazy_choices_completer(list_instances)

    try:
        system_paasta_config = load_system_paasta_config()
        valid_clusters = system_paasta_config.get_clusters()
        default_spark_cluster = system_paasta_config.get_spark_run_config().get(
            "default_cluster"
        )
        default_spark_pool = system_paasta_config.get_spark_run_config().get(
            "default_pool"
        )
    except PaastaNotConfiguredError:
        default_spark_cluster = "pnw-devc"
        default_spark_pool = "batch"
        valid_clusters = ["spark-pnw-prod", "pnw-devc"]

    list_parser.add_argument(
        "-c",
        "--cluster",
        help="The name of the cluster you wish to run Spark on.",
        choices=valid_clusters,
        default=default_spark_cluster,
    )

    list_parser.add_argument(
        "-p",
        "--pool",
        help="Name of the resource pool to run the Spark job.",
        default=default_spark_pool,
    )

    list_parser.add_argument(
        "-w",
        "--work-dir",
        default="{}:{}".format(os.getcwd(), DEFAULT_SPARK_WORK_DIR),
        help="The read-write volume to mount in format local_abs_dir:container_abs_dir",
    )

    list_parser.add_argument(
        "-y",
        "--yelpsoa-config-root",
        dest="yelpsoa_config_root",
        help="A directory from which yelpsoa-configs should be read from.",
        default=DEFAULT_SOA_DIR,
    )

    list_parser.add_argument(
        "-C",
        "--cmd",
        help="Run the spark-shell, pyspark, spark-submit, jupyter-lab, or history-server command.",
    )

    list_parser.add_argument(
        "--timeout-job-runtime",
        type=str,
        help="Timeout value which will be added before spark-submit. Job will exit if it doesn't "
        "finishes in given runtime. Recommended value: 2 * expected runtime. Example: 1h, 30m {DEFAULT_RUNTIME_TIMEOUT}",
        default=DEFAULT_RUNTIME_TIMEOUT,
    )

    list_parser.add_argument(
        "-d",
        "--dry-run",
        help="Shows the arguments supplied to docker as json.",
        action="store_true",
        default=False,
    )

    list_parser.add_argument(
        "--spark-args",
        help="Spark configurations documented in https://spark.apache.org/docs/latest/configuration.html, separated by space. "
        'For example, --spark-args "spark.executor.cores=1 spark.executor.memory=7g spark.executor.instances=2".',
    )

    list_parser.add_argument(
        "--nvidia",
        help="Use nvidia docker runtime for Spark driver process (requires GPU)",
        action="store_true",
        default=False,
    )

    list_parser.add_argument(
        "--mrjob",
        help="Pass Spark arguments to invoked command in the format expected by mrjobs",
        action="store_true",
        default=False,
    )

    list_parser.add_argument(
        "--cluster-manager",
        help="Specify which cluster manager to use. Support for certain cluster managers may be experimental",
        dest="cluster_manager",
        choices=CLUSTER_MANAGERS,
        default=CLUSTER_MANAGER_K8S,
    )

    list_parser.add_argument(
        "--enable-dra",
        help=(
            "[DEPRECATED] Enable Dynamic Resource Allocation (DRA) for the Spark job as documented in (y/spark-dra)."
            "DRA is enabled by default now. This config is a no-op operation and recommended to be removed."
        ),
        action="store_true",
        default=False,
    )

    k8s_target_cluster_type_group = list_parser.add_mutually_exclusive_group()
    k8s_target_cluster_type_group.add_argument(
        "--force-use-eks",
        help="Use the EKS version of the target cluster rather than the Yelp-managed target cluster",
        action="store_true",
        dest="use_eks_override",
        # We'll take a boolean value to mean that we should honor what the user wants, and None as using
        # the CI-provided default
        default=None,
    )
    k8s_target_cluster_type_group.add_argument(
        "--force-no-use-eks",
        help="Use the Yelp-managed version of the target cluster rather than the AWS-managed EKS target cluster",
        action="store_false",
        dest="use_eks_override",
        # We'll take a boolean value to mean that we should honor what the user wants, and None as using
        # the CI-provided default
        default=None,
    )

    list_parser.add_argument(
        "-j", "--jars", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    list_parser.add_argument(
        "--executor-memory", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    list_parser.add_argument(
        "--executor-cores", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    list_parser.add_argument(
        "--max-cores", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    list_parser.add_argument(
        "--driver-max-result-size", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    list_parser.add_argument(
        "--driver-memory", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    list_parser.add_argument(
        "--driver-cores", help=argparse.SUPPRESS, action=DeprecatedAction
    )

    aws_group = list_parser.add_argument_group(
        title="AWS credentials options",
        description="If --aws-credentials-yaml is specified, it overrides all "
        "other options. Otherwise, if -s/--service is specified, spark-run "
        "looks for service credentials in /etc/boto_cfg/[service].yaml. If "
        "it does not find the service credentials or no service is "
        "specified, spark-run falls back to the boto default behavior "
        "(checking ~/.aws/credentials, ~/.boto, etc).",
    )

    aws_group.add_argument(
        "--aws-credentials-yaml",
        help="Load aws keys from the provided yaml file. The yaml file must "
        "have keys for aws_access_key_id and aws_secret_access_key.",
    )

    aws_group.add_argument(
        "--aws-profile",
        help="Name of the AWS profile to load credentials from. Only used when "
        "--aws-credentials-yaml is not specified and --service is either "
        "not specified or the service does not have credentials in "
        "/etc/boto_cfg",
        default="default",
    )

    aws_group.add_argument(
        "--no-aws-credentials",
        help="Do not load any AWS credentials; allow the Spark job to use its "
        "own logic to load credentials",
        action="store_true",
        default=False,
    )

    aws_group.add_argument(
        "--aws-region",
        help=f"Specify an aws region. If the region is not specified, we will"
        f"default to using {DEFAULT_AWS_REGION}.",
        default=DEFAULT_AWS_REGION,
    )

    aws_group.add_argument(
        "--assume-aws-role",
        help="Takes an AWS IAM role ARN and attempts to create a session",
    )

    aws_group.add_argument(
        "--aws-role-duration",
        help=(
            "Duration in seconds for the role if --assume-aws-role provided. "
            "The maximum is 43200, but by default, roles may only allow 3600."
        ),
        type=int,
        default=43200,
    )

    jupyter_group = list_parser.add_argument_group(
        title="Jupyter kernel culling options",
        description="Idle kernels will be culled by default. Idle "
        "kernels with connections can be overridden not to be culled.",
    )

    jupyter_group.add_argument(
        "--cull-idle-timeout",
        type=int,
        default=7200,
        help="Timeout (in seconds) after which a kernel is considered idle and "
        "ready to be culled.",
    )

    jupyter_group.add_argument(
        "--not-cull-connected",
        action="store_true",
        default=False,
        help="By default, connected idle kernels are culled after timeout. "
        "They can be skipped if not-cull-connected is specified.",
    )

    list_parser.set_defaults(command=paasta_spark_run)


def decide_final_eks_toggle_state(user_override: Optional[bool]) -> bool:
    """
    This is slightly weird (hooray for tri-value logic!) - but basically:
    we want to prioritize any user choice for using EKS (i.e., force
    enable/disable using EKS) regardless of what the PaaSTA-supplied
    default is.

    If a user hasn't set --force-use-eks or --force-no-use-eks, argparse
    will leave args.use_eks_override (or, in this function, user_override) as None -
    otherwise, there'll be an actual boolean there.
    """
    if user_override is not None:
        return user_override

    return load_system_paasta_config().get_spark_use_eks_default()


def sanitize_container_name(container_name):
    # container_name only allows [a-zA-Z0-9][a-zA-Z0-9_.-]
    return re.sub("[^a-zA-Z0-9_.-]", "_", re.sub("^[^a-zA-Z0-9]+", "", container_name))


def generate_pod_template_path():
    return POD_TEMPLATE_PATH.format(file_uuid=uuid.uuid4().hex)


def should_enable_compact_bin_packing(disable_compact_bin_packing, cluster_manager):
    if disable_compact_bin_packing:
        return False

    if cluster_manager != CLUSTER_MANAGER_K8S:
        log.warn(
            "enable_compact_bin_packing=True ignored as cluster manager is not kubernetes"
        )
        return False

    if not os.access(POD_TEMPLATE_DIR, os.W_OK):
        log.warn(
            f"enable_compact_bin_packing=True ignored as {POD_TEMPLATE_DIR} is not usable"
        )
        return False

    return True


def get_docker_run_cmd(
    container_name,
    volumes,
    env,
    docker_img,
    docker_cmd,
    nvidia,
    docker_memory_limit,
    docker_shm_size,
    docker_cpu_limit,
):
    print(
        f"Setting docker memory, shared memory, and cpu limits as {docker_memory_limit}, {docker_shm_size}, and {docker_cpu_limit} core(s) respectively."
    )
    cmd = ["paasta_docker_wrapper", "run"]
    cmd.append(f"--memory={docker_memory_limit}")
    if docker_shm_size is not None:
        cmd.append(f"--shm-size={docker_shm_size}")
        cmd.append("--ulimit")
        cmd.append("memlock=-1")
    cmd.append(f"--cpus={docker_cpu_limit}")
    cmd.append("--rm")
    cmd.append("--net=host")

    non_interactive_cmd = ["spark-submit", "history-server"]
    if not any(c in docker_cmd for c in non_interactive_cmd):
        cmd.append("--interactive=true")
        if sys.stdout.isatty():
            cmd.append("--tty=true")

    cmd.append("--user=%d:%d" % (os.geteuid(), os.getegid()))
    cmd.append("--name=%s" % sanitize_container_name(container_name))
    for k, v in env.items():
        cmd.append("--env")
        if k in SENSITIVE_ENV:
            cmd.append(k)
        else:
            cmd.append(f"{k}={v}")
    if nvidia:
        cmd.append("--env")
        cmd.append("NVIDIA_VISIBLE_DEVICES=all")
        cmd.append("--runtime=nvidia")
    for volume in volumes:
        cmd.append("--volume=%s" % volume)
    cmd.append("%s" % docker_img)
    cmd.extend(("sh", "-c", docker_cmd))

    return cmd


def get_docker_image(
    args: argparse.Namespace, instance_config: InstanceConfig
) -> Optional[str]:
    """
    Since the Docker image digest used to launch the Spark cluster is obtained by inspecting local
    Docker images, we need to ensure that the Docker image exists locally or is pulled in all scenarios.
    """
    # docker image is built locally then pushed
    if args.build:
        return build_and_push_docker_image(args)

    docker_url = ""
    if args.image:
        docker_url = args.image
    else:
        try:
            docker_url = instance_config.get_docker_url()
        except NoDockerImageError:
            print(
                PaastaColors.red(
                    "Error: No sha has been marked for deployment for the %s deploy group.\n"
                    "Please ensure this service has either run through a jenkins pipeline "
                    "or paasta mark-for-deployment has been run for %s\n"
                    % (instance_config.get_deploy_group(), args.service)
                ),
                sep="",
                file=sys.stderr,
            )
            return None

    print(
        "Please wait while the image (%s) is pulled (times out after 5m)..."
        % docker_url,
        file=sys.stderr,
    )
    # Need sudo for credentials when pulling images from paasta docker registry (docker-paasta.yelpcorp.com)
    # However, in CI env, we can't connect to docker via root and we can pull with user `jenkins`
    is_ci_env = "CI" in os.environ
    cmd_prefix = "" if is_ci_env else "sudo -H "
    retcode, _ = _run(f"{cmd_prefix}docker pull {docker_url}", stream=True, timeout=300)
    if retcode != 0:
        print(
            "\nPull failed. Are you authorized to run docker commands?",
            file=sys.stderr,
        )
        return None
    return docker_url


def get_smart_paasta_instance_name(args):
    if os.environ.get("TRON_JOB_NAMESPACE"):
        tron_job = os.environ.get("TRON_JOB_NAME")
        tron_action = os.environ.get("TRON_ACTION")
        return f"{tron_job}.{tron_action}"
    else:
        how_submitted = None
        if args.mrjob:
            how_submitted = "mrjob"
        else:
            for spark_cmd in SPARK_COMMANDS:
                if spark_cmd in args.cmd:
                    how_submitted = spark_cmd
                    break
        how_submitted = how_submitted or "other"
        return f"{args.instance}_{get_username()}_{how_submitted}"


def get_spark_env(
    args: argparse.Namespace,
    spark_conf_str: str,
    aws_creds: Tuple[Optional[str], Optional[str], Optional[str]],
    ui_port: str,
    system_paasta_config: SystemPaastaConfig,
) -> Dict[str, str]:
    """Create the env config dict to configure on the docker container"""

    spark_env = {}
    access_key, secret_key, session_token = aws_creds
    if access_key:
        spark_env["AWS_ACCESS_KEY_ID"] = access_key
        spark_env["AWS_SECRET_ACCESS_KEY"] = secret_key
        if session_token is not None:
            spark_env["AWS_SESSION_TOKEN"] = session_token

    spark_env["AWS_DEFAULT_REGION"] = args.aws_region
    spark_env["PAASTA_LAUNCHED_BY"] = get_possible_launched_by_user_variable_from_env()
    spark_env["PAASTA_INSTANCE_TYPE"] = "spark"

    # Run spark (and mesos framework) as root.
    spark_env["SPARK_USER"] = "root"
    spark_env["SPARK_OPTS"] = spark_conf_str

    # Default configs to start the jupyter notebook server
    if args.cmd == "jupyter-lab":
        spark_env["JUPYTER_RUNTIME_DIR"] = "/source/.jupyter"
        spark_env["JUPYTER_DATA_DIR"] = "/source/.jupyter"
        spark_env["JUPYTER_CONFIG_DIR"] = "/source/.jupyter"
    elif args.cmd == "history-server":
        dirs = args.work_dir.split(":")
        spark_env["SPARK_LOG_DIR"] = dirs[1]
        if not args.spark_args or not args.spark_args.startswith(
            "spark.history.fs.logDirectory"
        ):
            print(
                "history-server requires spark.history.fs.logDirectory in spark-args",
                file=sys.stderr,
            )
            sys.exit(1)
        spark_env["SPARK_HISTORY_OPTS"] = (
            f"-D{args.spark_args} " f"-Dspark.history.ui.port={ui_port}"
        )
        spark_env["SPARK_DAEMON_CLASSPATH"] = "/opt/spark/extra_jars/*"
        spark_env["SPARK_NO_DAEMONIZE"] = "true"

    if decide_final_eks_toggle_state(args.use_eks_override):
        spark_env["KUBECONFIG"] = system_paasta_config.get_spark_kubeconfig()

    return spark_env


def _parse_user_spark_args(
    spark_args: Optional[str],
    pod_template_path: str,
    enable_compact_bin_packing: bool = False,
    enable_spark_dra: bool = False,
) -> Dict[str, str]:

    user_spark_opts = {}
    if spark_args:
        for spark_arg in spark_args.split():
            fields = spark_arg.split("=", 1)
            if len(fields) != 2:
                print(
                    PaastaColors.red(
                        "Spark option %s is not in format option=value." % spark_arg
                    ),
                    file=sys.stderr,
                )
                sys.exit(1)
            user_spark_opts[fields[0]] = fields[1]

    if enable_compact_bin_packing:
        user_spark_opts["spark.kubernetes.executor.podTemplateFile"] = pod_template_path

    if enable_spark_dra:
        if (
            "spark.dynamicAllocation.enabled" in user_spark_opts
            and user_spark_opts["spark.dynamicAllocation.enabled"] == "false"
        ):
            print(
                PaastaColors.red(
                    "Error: --enable-dra flag is provided while spark.dynamicAllocation.enabled "
                    "is explicitly set to false in --spark-args. If you want to enable DRA, please remove the "
                    "spark.dynamicAllocation.enabled=false config from spark-args. If you don't want to "
                    "enable DRA, please remove the --enable-dra flag."
                ),
                file=sys.stderr,
            )
            sys.exit(1)
        user_spark_opts["spark.dynamicAllocation.enabled"] = "true"

    return user_spark_opts


def create_spark_config_str(spark_config_dict, is_mrjob):
    conf_option = "--jobconf" if is_mrjob else "--conf"
    spark_config_entries = list()

    if is_mrjob:
        spark_master = spark_config_dict["spark.master"]
        spark_config_entries.append(f"--spark-master={spark_master}")

    for opt, val in spark_config_dict.items():
        # mrjob use separate options to configure master
        if is_mrjob and opt == "spark.master":
            continue
        # Process Spark configs with multiple space separated values to be in single quotes
        if isinstance(val, str) and " " in val:
            val = f"'{val}'"
        spark_config_entries.append(f"{conf_option} {opt}={val}")
    return " ".join(spark_config_entries)


def run_docker_container(
    container_name,
    volumes,
    environment,
    docker_img,
    docker_cmd,
    dry_run,
    nvidia,
    docker_memory_limit,
    docker_shm_size,
    docker_cpu_limit,
) -> int:

    docker_run_args = dict(
        container_name=container_name,
        volumes=volumes,
        env=environment,
        docker_img=docker_img,
        docker_cmd=docker_cmd,
        nvidia=nvidia,
        docker_memory_limit=docker_memory_limit,
        docker_shm_size=docker_shm_size,
        docker_cpu_limit=docker_cpu_limit,
    )
    docker_run_cmd = get_docker_run_cmd(**docker_run_args)
    if dry_run:
        print(json.dumps(docker_run_cmd))
        return 0

    merged_env = {**os.environ, **environment}
    os.execlpe("paasta_docker_wrapper", *docker_run_cmd, merged_env)
    return 0


def get_spark_app_name(original_docker_cmd: Union[Any, str, List[str]]) -> str:
    """Use submitted batch name as default spark_run job name"""
    docker_cmds = (
        shlex.split(original_docker_cmd)
        if isinstance(original_docker_cmd, str)
        else original_docker_cmd
    )
    spark_app_name = None
    after_spark_submit = False
    for arg in docker_cmds:
        if arg == "spark-submit":
            after_spark_submit = True
        elif after_spark_submit and arg.endswith(".py"):
            batch_name = arg.split("/")[-1].replace(".py", "")
            spark_app_name = "paasta_" + batch_name
            break
        elif arg == "jupyter-lab":
            spark_app_name = "paasta_jupyter"
            break

    if spark_app_name is None:
        spark_app_name = "paasta_spark_run"

    spark_app_name += f"_{get_username()}"

    return spark_app_name


def _calculate_docker_memory_limit(
    spark_conf: Mapping[str, str], memory_limit: Optional[str]
) -> str:
    """In Order of preference:
    1. Argument: --docker-memory-limit
    2. --spark-args or spark-submit: spark.driver.memory
    3. Default
    """
    if memory_limit:
        return memory_limit

    try:
        docker_memory_limit_str = spark_conf.get(
            "spark.driver.memory", DEFAULT_DRIVER_MEMORY_BY_SPARK
        )
        adjustment_factor = DOCKER_RESOURCE_ADJUSTMENT_FACTOR
        match = re.match(r"([0-9]+)([a-z]*)", docker_memory_limit_str)
        memory_val = int(match[1]) * adjustment_factor
        memory_unit = match[2]
        docker_memory_limit = f"{memory_val}{memory_unit}"
    except Exception as e:
        # For any reason it fails, continue with default value
        print(
            f"ERROR: Failed to parse docker memory limit. Error: {e}. Example values: 1g, 200m."
        )
        raise

    return docker_memory_limit


def _calculate_docker_shared_memory_size(shm_size: Optional[str]) -> str:
    """In Order of preference:
    1. Argument: --docker-shm-size
    3. Default
    """
    if shm_size:
        return shm_size

    return DEFAULT_DOCKER_SHM_SIZE


def _calculate_docker_cpu_limit(
    spark_conf: Mapping[str, str], cpu_limit: Optional[str]
) -> str:
    """In Order of preference:
    1. Argument: --docker-cpu-limit
    2. --spark-args or spark-submit: spark.driver.cores
    3. Default
    """
    return (
        cpu_limit
        if cpu_limit
        else spark_conf.get("spark.driver.cores", str(DEFAULT_DRIVER_CORES_BY_SPARK))
    )


def configure_and_run_docker_container(
    args: argparse.Namespace,
    docker_img: str,
    instance_config: InstanceConfig,
    system_paasta_config: SystemPaastaConfig,
    spark_conf: Mapping[str, str],
    aws_creds: Tuple[Optional[str], Optional[str], Optional[str]],
    cluster_manager: str,
    pod_template_path: str,
) -> int:
    docker_memory_limit = _calculate_docker_memory_limit(
        spark_conf, args.docker_memory_limit
    )
    docker_shm_size = _calculate_docker_shared_memory_size(args.docker_shm_size)
    docker_cpu_limit = _calculate_docker_cpu_limit(
        spark_conf,
        args.docker_cpu_limit,
    )

    if cluster_manager in {CLUSTER_MANAGER_K8S, CLUSTER_MANAGER_LOCAL}:
        # service_configuration_lib puts volumes into the k8s
        # configs for local mode
        volumes = get_volumes_from_spark_k8s_configs(spark_conf)
    else:
        raise UnsupportedClusterManagerException(cluster_manager)

    volumes.append("%s:rw" % args.work_dir)
    volumes.append("/nail/home:/nail/home:rw")

    if args.enable_compact_bin_packing:
        volumes.append(f"{pod_template_path}:{pod_template_path}:rw")

    if decide_final_eks_toggle_state(args.use_eks_override):
        volumes.append(
            f"{system_paasta_config.get_spark_kubeconfig()}:{system_paasta_config.get_spark_kubeconfig()}:ro"
        )

    environment = instance_config.get_env_dictionary()  # type: ignore
    spark_conf_str = create_spark_config_str(spark_conf, is_mrjob=args.mrjob)
    environment.update(
        get_spark_env(
            args=args,
            spark_conf_str=spark_conf_str,
            aws_creds=aws_creds,
            ui_port=spark_conf["spark.ui.port"],
            system_paasta_config=system_paasta_config,
        )
    )  # type:ignore

    webui_url = get_webui_url(spark_conf["spark.ui.port"])
    webui_url_msg = PaastaColors.green(f"\nSpark monitoring URL: ") + f"{webui_url}\n"

    docker_cmd = get_docker_cmd(args, instance_config, spark_conf_str)
    if "history-server" in docker_cmd:
        print(PaastaColors.green(f"\nSpark history server URL: ") + f"{webui_url}\n")
    elif any(c in docker_cmd for c in ["pyspark", "spark-shell", "spark-submit"]):
        grafana_url = get_grafana_url(spark_conf)
        signalfx_url = get_signalfx_url(spark_conf)
        dashboard_url_msg = (
            PaastaColors.green(f"\nGrafana dashboard: ")
            + f"{grafana_url}\n"
            + PaastaColors.green(f"\nSignalfx dashboard: ")
            + f"{signalfx_url}\n"
        )
        print(webui_url_msg)
        print(dashboard_url_msg)
        log.info(webui_url_msg)
        log.info(dashboard_url_msg)
        spark_conf_builder = spark_config.SparkConfBuilder()
        history_server_url = spark_conf_builder.get_history_url(spark_conf)
        if history_server_url:
            history_server_url_msg = (
                f"\nAfter the job is finished, you can find the spark UI from {history_server_url}\n"
                "Check y/spark-recent-history for faster access to prod logs\n"
            )
            print(history_server_url_msg)
            log.info(history_server_url_msg)
    print(f"Selected cluster manager: {cluster_manager}\n")

    if clusterman_metrics and _should_get_resource_requirements(docker_cmd, args.mrjob):
        resources = get_resources_requested(spark_conf)
        hourly_cost = get_spark_hourly_cost(
            clusterman_metrics,
            resources,
            spark_conf["spark.executorEnv.PAASTA_CLUSTER"],
            args.pool,
        )
        message = (
            f"Resource request ({resources['cpus']} cpus and {resources['mem']} MB memory total)"
            f" is estimated to cost ${hourly_cost} per hour"
        )
        if clusterman_metrics.util.costs.should_warn(hourly_cost):
            print(PaastaColors.red(f"WARNING: {message}"))
        else:
            print(message)

    return run_docker_container(
        container_name=spark_conf["spark.app.name"],
        volumes=volumes,
        environment=environment,
        docker_img=docker_img,
        docker_cmd=docker_cmd,
        dry_run=args.dry_run,
        nvidia=args.nvidia,
        docker_memory_limit=docker_memory_limit,
        docker_shm_size=docker_shm_size,
        docker_cpu_limit=docker_cpu_limit,
    )


def _should_get_resource_requirements(docker_cmd: str, is_mrjob: bool) -> bool:
    return is_mrjob or any(
        c in docker_cmd for c in ["pyspark", "spark-shell", "spark-submit"]
    )


def get_docker_cmd(
    args: argparse.Namespace, instance_config: InstanceConfig, spark_conf_str: str
) -> str:
    original_docker_cmd = str(args.cmd or instance_config.get_cmd())

    if args.mrjob:
        return original_docker_cmd + " " + spark_conf_str
    # Default cli options to start the jupyter notebook server.
    elif original_docker_cmd == "jupyter-lab":
        cull_opts = (
            "--MappingKernelManager.cull_idle_timeout=%s " % args.cull_idle_timeout
        )
        if args.not_cull_connected is False:
            cull_opts += "--MappingKernelManager.cull_connected=True "

        return "SHELL=bash USER={} /source/virtualenv_run_jupyter/bin/jupyter-lab -y --ip={} {}".format(
            get_username(), socket.getfqdn(), cull_opts
        )
    elif original_docker_cmd == "history-server":
        return "start-history-server.sh"
    # Spark options are passed as options to pyspark and spark-shell.
    # For jupyter, environment variable SPARK_OPTS is set instead.
    else:
        return inject_spark_conf_str(original_docker_cmd, spark_conf_str)


def _get_adhoc_docker_registry(service: str, soa_dir: str = DEFAULT_SOA_DIR) -> str:
    if service is None:
        raise NotImplementedError('"None" is not a valid service')

    service_configuration = read_service_configuration(service, soa_dir)
    return service_configuration.get("docker_registry", DEFAULT_SPARK_DOCKER_REGISTRY)


def build_and_push_docker_image(args: argparse.Namespace) -> Optional[str]:
    """
    Build an image if the default Spark service image is not preferred.
    The image needs to be pushed to a registry for the Spark executors
    to pull.
    """
    if not makefile_responds_to("cook-image"):
        print(
            "A local Makefile with a 'cook-image' target is required for --build",
            file=sys.stderr,
        )
        return None

    default_tag = "{}-{}".format(DEFAULT_SPARK_DOCKER_IMAGE_PREFIX, get_username())
    docker_tag = os.environ.get("DOCKER_TAG", default_tag)
    os.environ["DOCKER_TAG"] = docker_tag

    cook_return = paasta_cook_image(
        args=None, service=args.service, soa_dir=args.yelpsoa_config_root
    )
    if cook_return != 0:
        return None

    registry_uri = args.docker_registry or _get_adhoc_docker_registry(
        service=args.service,
        soa_dir=args.yelpsoa_config_root,
    )

    docker_url = f"{registry_uri}/{docker_tag}"
    command = f"docker tag {docker_tag} {docker_url}"
    print(PaastaColors.grey(command))
    retcode, _ = _run(command, stream=True)
    if retcode != 0:
        return None

    if registry_uri != DEFAULT_SPARK_DOCKER_REGISTRY:
        command = "sudo -H docker push %s" % docker_url
    else:
        command = "docker push %s" % docker_url

    print(PaastaColors.grey(command))
    retcode, output = _run(command, stream=False)
    if retcode != 0:
        return None

    # With unprivileged docker, the digest on the remote registry may not match the digest
    # in the local environment. Because of this, we have to parse the digest message from the
    # server response and use downstream when launching spark executors

    # Output from `docker push` with unprivileged docker looks like
    #  Using default tag: latest
    #  The push refers to repository [docker-dev.yelpcorp.com/paasta-spark-run-dpopes:latest]
    #  latest: digest: sha256:0a43aa65174a400bd280d48d460b73eb49b0ded4072c9e173f919543bf693557

    # With privileged docker, the last line has an extra "size: 123"
    #  latest: digest: sha256:0a43aa65174a400bd280d48d460b73eb49b0ded4072c9e173f919543bf693557 size: 52

    digest_line = output.split("\n")[-1]
    digest_match = re.match(r"[^:]*: [^:]*: (?P<digest>[^\s]*)", digest_line)
    if not digest_match:
        raise ValueError(f"Could not determine digest from output: {output}")
    digest = digest_match.group("digest")

    image_url = f"{docker_url}@{digest}"

    # If the local digest doesn't match the remote digest AND the registry is
    # non-default (which requires requires authentication, and consequently sudo),
    # downstream `docker run` commands will fail trying to authenticate.
    # To work around this, we can proactively `sudo docker pull` here so that
    # the image exists locally and can be `docker run` without sudo
    if registry_uri != DEFAULT_SPARK_DOCKER_REGISTRY:
        command = f"sudo -H docker pull {image_url}"
        print(PaastaColors.grey(command))
        retcode, output = _run(command, stream=False)
        if retcode != 0:
            raise NoDockerImageError(f"Could not pull {image_url}: {output}")

    return image_url


def validate_work_dir(s):
    dirs = s.split(":")
    if len(dirs) != 2:
        print(
            "work-dir %s is not in format local_abs_dir:container_abs_dir" % s,
            file=sys.stderr,
        )
        sys.exit(1)

    for d in dirs:
        if not os.path.isabs(d):
            print("%s is not an absolute path" % d, file=sys.stderr)
            sys.exit(1)


def _auto_add_timeout_for_job(cmd, timeout_job_runtime):
    # Timeout only to be added for spark-submit commands
    # TODO: Add timeout for jobs using mrjob with spark-runner
    if "spark-submit" not in cmd:
        return cmd
    try:
        timeout_present = re.match(
            r"^.*timeout[\s]+[\d]+[\.]?[\d]*[m|h][\s]+spark-submit .*$", cmd
        )
        if not timeout_present:
            split_cmd = cmd.split("spark-submit")
            cmd = f"{split_cmd[0]}timeout {timeout_job_runtime} spark-submit{split_cmd[1]}"
            print(
                PaastaColors.blue(
                    f"NOTE: Job will exit in given time {timeout_job_runtime}. "
                    f"Adjust timeout value using --timeout-job-timeout. "
                    f"New Updated Command with timeout: {cmd}"
                ),
            )
    except Exception as e:
        err_msg = (
            f"'timeout' could not be added to command: '{cmd}' due to error '{e}'. "
            "Please report to #spark."
        )
        log.warn(err_msg)
        print(PaastaColors.red(err_msg))
    return cmd


def _validate_pool(args, system_paasta_config):
    if args.pool:
        valid_pools = system_paasta_config.get_cluster_pools().get(args.cluster, [])
        if not valid_pools:
            log.warning(
                PaastaColors.yellow(
                    f"Could not fetch allowed_pools for `{args.cluster}`. Skipping pool validation.\n"
                )
            )
        if valid_pools and args.pool not in valid_pools:
            print(
                PaastaColors.red(
                    f"Invalid --pool value. List of valid pools for cluster `{args.cluster}`: "
                    f"{valid_pools}"
                ),
                file=sys.stderr,
            )
            return False
    return True


def _get_k8s_url_for_cluster(cluster: str) -> Optional[str]:
    """
    Annoyingly, there's two layers of aliases: one to figure out what
    k8s server url to use (this one) and another to figure out what
    soaconfigs filename to use ;_;

    This exists so that we can map something like `--cluster pnw-devc`
    into spark-pnw-devc's k8s apiserver url without needing to update
    any soaconfigs/alter folk's muscle memory.

    Ideally we can get rid of this entirely once spark-run reads soaconfigs
    in a manner more closely aligned to what we do with other paasta workloads
    (i.e., have it automatically determine where to run based on soaconfigs
    filenames - and not rely on explicit config)
    """
    realized_cluster = (
        load_system_paasta_config().get_eks_cluster_aliases().get(cluster, cluster)
    )
    return (
        load_system_paasta_config()
        .get_kube_clusters()
        .get(realized_cluster, {})
        .get("server")
    )


def paasta_spark_run(args):
    # argparse does not work as expected with both default and
    # type=validate_work_dir.
    validate_work_dir(args.work_dir)

    try:
        system_paasta_config = load_system_paasta_config()
    except PaastaNotConfiguredError:
        print(
            PaastaColors.yellow(
                "Warning: Couldn't load config files from '/etc/paasta'. This indicates"
                "PaaSTA is not configured locally on this host, and local-run may not behave"
                "the same way it would behave on a server configured for PaaSTA."
            ),
            sep="\n",
        )
        system_paasta_config = SystemPaastaConfig({"volumes": []}, "/etc/paasta")

    if args.cmd == "jupyter-lab" and not args.build and not args.image:
        print(
            PaastaColors.red(
                "The jupyter-lab command requires a prebuilt image with -I or --image."
            ),
            file=sys.stderr,
        )
        return 1

    # validate pool
    if not _validate_pool(args, system_paasta_config):
        return 1

    # annoyingly, there's two layers of aliases: one for the soaconfigs to read from
    # (that's this alias lookup) - and then another layer later when figuring out what
    # k8s server url to use ;_;
    cluster = system_paasta_config.get_cluster_aliases().get(args.cluster, args.cluster)
    # Use the default spark:client instance configs if not provided
    try:
        instance_config = get_instance_config(
            service=args.service,
            instance=args.instance,
            cluster=cluster,
            load_deployments=args.build is False and args.image is None,
            soa_dir=args.yelpsoa_config_root,
        )
    except NoConfigurationForServiceError as e:
        print(str(e), file=sys.stderr)
        return 1
    except NoDeploymentsAvailable:
        print(
            PaastaColors.red(
                "Error: No deployments.json found in %(soa_dir)s/%(service)s."
                "You can generate this by running:"
                "generate_deployments_for_service -d %(soa_dir)s -s %(service)s"
                % {"soa_dir": args.yelpsoa_config_root, "service": args.service}
            ),
            sep="\n",
            file=sys.stderr,
        )
        return 1

    if not args.cmd and not instance_config.get_cmd():
        print(
            "A command is required, pyspark, spark-shell, spark-submit or jupyter",
            file=sys.stderr,
        )
        return 1

    aws_creds = get_aws_credentials(
        service=args.service,
        no_aws_credentials=args.no_aws_credentials,
        aws_credentials_yaml=args.aws_credentials_yaml,
        profile_name=args.aws_profile,
        assume_aws_role_arn=args.assume_aws_role,
        session_duration=args.aws_role_duration,
    )
    docker_image_digest = get_docker_image(args, instance_config)
    if docker_image_digest is None:
        return 1

    pod_template_path = generate_pod_template_path()
    args.enable_compact_bin_packing = should_enable_compact_bin_packing(
        args.disable_compact_bin_packing, args.cluster_manager
    )

    volumes = instance_config.get_volumes(system_paasta_config.get_volumes())
    app_base_name = get_spark_app_name(args.cmd or instance_config.get_cmd())

    if args.enable_compact_bin_packing:
        document = POD_TEMPLATE.format(
            spark_pod_label=limit_size_with_hash(f"exec-{app_base_name}"),
        )
        parsed_pod_template = yaml.safe_load(document)
        with open(pod_template_path, "w") as f:
            yaml.dump(parsed_pod_template, f)

    user_spark_opts = _parse_user_spark_args(
        args.spark_args,
        pod_template_path,
        args.enable_compact_bin_packing,
        args.enable_dra,
    )

    args.cmd = _auto_add_timeout_for_job(args.cmd, args.timeout_job_runtime)

    # This is required if configs are provided as part of `spark-submit`
    # Other way to provide is with --spark-args
    sub_cmds = args.cmd.split(" ")  # spark.driver.memory=10g
    for cmd in sub_cmds:
        if cmd.startswith("spark.driver.memory") or cmd.startswith(
            "spark.driver.cores"
        ):
            key, value = cmd.split("=")
            user_spark_opts[key] = value

    paasta_instance = get_smart_paasta_instance_name(args)

    use_eks = decide_final_eks_toggle_state(args.use_eks_override)
    k8s_server_address = _get_k8s_url_for_cluster(args.cluster) if use_eks else None
    paasta_cluster = (
        args.cluster
        if not use_eks
        else load_system_paasta_config()
        .get_eks_cluster_aliases()
        .get(args.cluster, args.cluster)
    )
    spark_conf_builder = spark_config.SparkConfBuilder()
    spark_conf = spark_conf_builder.get_spark_conf(
        cluster_manager=args.cluster_manager,
        spark_app_base_name=app_base_name,
        docker_img=docker_image_digest,
        user_spark_opts=user_spark_opts,
        paasta_cluster=paasta_cluster,
        paasta_pool=args.pool,
        paasta_service=args.service,
        paasta_instance=paasta_instance,
        extra_volumes=volumes,
        aws_creds=aws_creds,
        aws_region=args.aws_region,
        force_spark_resource_configs=args.force_spark_resource_configs,
        use_eks=use_eks,
        k8s_server_address=k8s_server_address,
    )
    # TODO: Remove this after MLCOMPUTE-699 is complete
    if "spark.kubernetes.decommission.script" not in spark_conf:
        spark_conf[
            "spark.kubernetes.decommission.script"
        ] = "/opt/spark/kubernetes/dockerfiles/spark/decom.sh"

    return configure_and_run_docker_container(
        args,
        docker_img=docker_image_digest,
        instance_config=instance_config,
        system_paasta_config=system_paasta_config,
        spark_conf=spark_conf,
        aws_creds=aws_creds,
        cluster_manager=args.cluster_manager,
        pod_template_path=pod_template_path,
    )
