"""Utilities for e2e tests of KubeRay/Ray integration.
For consistency, all K8s interactions use kubectl through subprocess calls.
"""
import atexit
import contextlib
import logging
import os
import pathlib
import subprocess
import tempfile
import time
from typing import Any, Dict, Generator, List, Optional

import yaml

logger = logging.getLogger(__name__)

SCRIPTS_DIR = pathlib.Path(__file__).resolve().parent / "scripts"
TEST_CR_PATH = (
    pathlib.Path(__file__).resolve().parent / "setup" / "raycluster_test.yaml"
)
TEST_CLUSTER_NAME = "raycluster-test"

# Parent directory of Ray repository
RAY_PARENT = str(pathlib.Path(__file__).resolve().parents[5])

RAYCLUSTERS_QUALIFIED = "rayclusters.ray.io"

LOG_FORMAT = "[%(levelname)s %(asctime)s] %(filename)s: %(lineno)d  %(message)s"


def switch_to_ray_parent_dir():
    # Switch to parent of Ray repo, because that's what the doc examples do.
    logger.info("Switching to parent of Ray directory.")
    os.chdir(RAY_PARENT)


def setup_kuberay_operator():
    """Set up KubeRay operator and Ray autoscaler RBAC."""
    switch_to_ray_parent_dir()
    logger.info("Cloning KubeRay and setting up KubeRay configuration.")
    # For faster run-time when triggering the test locally, don't run the init
    # script if it has already been run.
    subprocess.check_call(
        [
            "bash",
            "-c",
            (
                "ls ray/python/ray/autoscaler/kuberay/config ||"
                " ./ray/python/ray/autoscaler/kuberay/init-config.sh"
            ),
        ]
    )
    logger.info("Creating KubeRay operator.")
    subprocess.check_call(
        [
            "kubectl",
            "create",
            "-k",
            "ray/python/ray/autoscaler/kuberay/config/default",
        ]
    )


def teardown_kuberay_operator():
    logger.info("Switching to parent of Ray directory.")
    os.chdir(RAY_PARENT)

    logger.info("Deleting operator.")
    subprocess.check_call(
        [
            "kubectl",
            "delete",
            "--ignore-not-found",
            "-k",
            "ray/python/ray/autoscaler/kuberay/config/default",
        ]
    )

    logger.info("Double-checking no pods left over in namespace ray-system.")
    wait_for_pods(goal_num_pods=0, namespace="ray-system")


def wait_for_raycluster_crd(tries=60, backoff_s=5):
    """CRD creation can take a bit of time after the client request.
    This function waits until the crd with the provided name is registered.
    """
    switch_to_ray_parent_dir()
    logger.info("Making sure RayCluster CRD has been registered.")
    for i in range(tries):
        get_crd_output = subprocess.check_output(["kubectl", "get", "crd"]).decode()
        if RAYCLUSTERS_QUALIFIED in get_crd_output:
            logger.info("Confirmed existence of RayCluster CRD.")
            break
        elif i < tries - 1:
            logger.info("Still waiting to register RayCluster CRD.")
            time.sleep(backoff_s)
        else:
            raise Exception("Failed to register RayCluster CRD.")

    # Create a test RayCluster CR to make sure that the CRD is fully registered.
    for i in range(tries):
        try:
            subprocess.check_call(["kubectl", "apply", "-f", TEST_CR_PATH])
            break
        except subprocess.CalledProcessError as e:
            logger.info("Can't create RayCluster CR.")
            if i < tries - 1:
                logger.info("Retrying.")
                time.sleep(backoff_s)
            else:
                logger.info("Giving up.")
                raise e from None

    # Confirm the test RayCluster exists.
    out = subprocess.check_output(["kubectl", "get", RAYCLUSTERS_QUALIFIED]).decode()
    assert TEST_CLUSTER_NAME in out, out

    # Delete the test RayCluster.
    subprocess.check_call(["kubectl", "delete", "-f", TEST_CR_PATH])
    # Make sure the associated resources are gone before proceeding.
    wait_for_pods(goal_num_pods=0, namespace="default")


def wait_for_pods(goal_num_pods: int, namespace: str, tries=60, backoff_s=5) -> None:
    """Wait for the number of pods in the `namespace` to be exactly `num_pods`.

    Raise an exception after exceeding `tries` attempts with `backoff_s` second waits.
    """
    for i in range(tries):
        cur_num_pods = _get_num_pods(namespace)
        if cur_num_pods == goal_num_pods:
            logger.info(f"Confirmed {goal_num_pods} pod(s) in namespace {namespace}.")
            return
        elif i < tries - 1:
            logger.info(
                f"The number of pods in namespace {namespace} is {cur_num_pods}."
                f" Waiting until the number of pods is {goal_num_pods}."
                f"{get_pod_names(namespace)}"
            )
            time.sleep(backoff_s)
        else:
            raise Exception(
                f"Failed to scale to {goal_num_pods} pod(s) in namespace {namespace}."
            )


def _get_num_pods(namespace: str) -> int:
    return len(get_pod_names(namespace))


def get_pod_names(namespace: str) -> List[str]:
    """Get the list of pod names in the namespace."""
    get_pods_output = (
        subprocess.check_output(
            [
                "kubectl",
                "-n",
                namespace,
                "get",
                "pods",
                "-o",
                "custom-columns=POD:metadata.name",
                "--no-headers",
            ]
        )
        .decode()
        .strip()
    )

    # If there aren't any pods, the output is any empty string.
    if not get_pods_output:
        return []
    else:
        return get_pods_output.split("\n")


def wait_for_pod_to_start(
    pod_name_filter: str, namespace: str, tries=60, backoff_s=5
) -> None:
    """Waits for a pod to have Running status.phase.

    More precisely, waits until there is a pod with name containing `pod_name_filter`
    and the pod has Running status.phase."""
    for i in range(tries):
        pod = get_pod(pod_name_filter=pod_name_filter, namespace=namespace)
        if not pod:
            # We didn't get a matching pod.
            continue
        pod_status = (
            subprocess.check_output(
                [
                    "kubectl",
                    "-n",
                    namespace,
                    "get",
                    "pod",
                    pod,
                    "-o",
                    "custom-columns=POD:status.phase",
                    "--no-headers",
                ]
            )
            .decode()
            .strip()
        )
        # "not found" is part of the kubectl output if the pod's not there.
        if "not found" in pod_status:
            raise Exception(f"Pod {pod} not found.")
        elif pod_status == "Running":
            logger.info(f"Confirmed pod {pod} is Running.")
            return
        elif i < tries - 1:
            logger.info(
                f"Pod {pod} has status {pod_status}. Waiting for the pod to enter "
                "Running status."
            )
            time.sleep(backoff_s)
        else:
            raise Exception(f"Timed out waiting for pod {pod} to enter Running status.")


def wait_for_ray_health(
    pod_name_filter: str,
    namespace: str,
    tries=60,
    backoff_s=5,
    ray_container="ray-head",
) -> None:
    """Waits until a Ray pod passes `ray health-check`.

    More precisely, waits until a Ray pod whose name includes the string
    `pod_name_filter` passes `ray health-check`.
    (Ensures Ray has completely started in the pod.)

    Use case: Wait until there is a Ray head pod with Ray running on it.
    """
    for i in range(tries):
        try:
            pod = get_pod(pod_name_filter=pod_name_filter, namespace="default")
            assert pod, f"Couldn't find a pod matching {pod_name_filter}."
            # `ray health-check` yields 0 exit status iff it succeeds
            kubectl_exec(
                ["ray", "health-check"], pod, namespace, container=ray_container
            )
            logger.info(f"ray health check passes for pod {pod}")
            return
        except subprocess.CalledProcessError as e:
            logger.info(f"Failed ray health check for pod {pod}.")
            if i < tries - 1:
                logger.info("Trying again.")
                time.sleep(backoff_s)
            else:
                logger.info("Giving up.")
                raise e from None


def get_pod(pod_name_filter: str, namespace: str) -> Optional[str]:
    """Gets pods in the `namespace`.

    Returns the first pod that has `pod_name_filter` as a
    substring of its name. Returns None if there are no matches.
    """
    pod_names = get_pod_names(namespace)
    matches = [pod_name for pod_name in pod_names if pod_name_filter in pod_name]
    if not matches:
        logger.warning(f"No match for `{pod_name_filter}` in namespace `{namespace}`.")
        return None
    return matches[0]


def kubectl_exec(
    command: List[str],
    pod: str,
    namespace: str,
    container: Optional[str] = None,
) -> str:
    """kubectl exec the `command` in the given `pod` in the given `namespace`.
    If a `container` is specified, will specify that container for kubectl.

    Prints and return kubectl's output as a string.
    """
    container_option = ["-c", container] if container else []
    kubectl_exec_command = (
        ["kubectl", "exec", "-it", pod] + container_option + ["--"] + command
    )
    # Print for debugging convenience.
    try:
        out = subprocess.check_output(kubectl_exec_command).decode().strip()
    except subprocess.CalledProcessError as e:
        logger.error(f"Error running command {kubectl_exec_command}.")
        logger.error(f"Output: {e.output.decode()}")
        raise e from None
    print(out)
    return out


def kubectl_logs(
    pod: str,
    namespace: str,
    container: Optional[str] = None,
) -> str:
    """Wrapper for kubectl logs.

    Returns the logs as a string.
    """
    container_option = ["-c", container] if container else []
    kubectl_logs_command = ["kubectl", "logs", pod] + container_option
    out = subprocess.check_output(kubectl_logs_command).decode().strip()
    return out


def kubectl_exec_python_script(
    script_name: str,
    pod: str,
    namespace: str,
    container: Optional[str] = None,
) -> str:
    """
    Runs a python script in a container via `kubectl exec`.
    Scripts live in `tests/kuberay/scripts`.

    Prints and return kubectl's output as a string.
    """
    script_path = SCRIPTS_DIR / script_name
    with open(script_path) as script_file:
        script_string = script_file.read()
    return kubectl_exec(["python", "-c", script_string], pod, namespace, container)


def get_raycluster(raycluster: str, namespace: str) -> Dict[str, Any]:
    """Gets the Ray CR with name `raycluster` in namespace `namespace`.

    Returns the CR as a nested Dict.
    """
    get_raycluster_output = (
        subprocess.check_output(
            ["kubectl", "-n", namespace, "get", "raycluster", raycluster, "-o", "yaml"]
        )
        .decode()
        .strip()
    )
    return yaml.safe_load(get_raycluster_output)


def _get_service_port(service: str, namespace: str, target_port: int) -> int:
    """Given a K8s service and a port targetted by the service, returns the
    corresponding port exposed by the service.

    Args:
        service: Name of a K8s service.
        namespace: Namespace to which the service belongs.
        target_port: Port targeted by the service.

    Returns:
        service_port: The port exposed by the service.
    """
    service_str = (
        subprocess.check_output(
            ["kubectl", "-n", namespace, "get", "service", service, "-o", "yaml"]
        )
        .decode()
        .strip()
    )
    service_dict = yaml.safe_load(service_str)
    service_ports: List = service_dict["spec"]["ports"]
    matching_ports = [
        port for port in service_ports if port["targetPort"] == target_port
    ]
    assert matching_ports
    service_port = matching_ports[0]["port"]
    return service_port


@contextlib.contextmanager
def _kubectl_port_forward(
    service: str, namespace: str, target_port: int, local_port: Optional[int] = None
) -> Generator[int, None, None]:
    """Context manager which creates a kubectl port-forward process targeting a
    K8s service.

    Terminates the port-forwarding process upon exit.

    Args:
        service: Name of a K8s service.
        namespace: Namespace to which the service belongs.
        target_port: The port targeted by the service.
        local_port: Forward from this port. Optional. By default, uses the port exposed
            by the service.

    Yields:
        int: The local port. The service can then be accessed at
            127.0.0.1:<local_port>.
    """
    # First, figure out which port the service exposes for the given target port.
    service_port = _get_service_port(service, namespace, target_port)
    if not local_port:
        local_port = service_port

    process = subprocess.Popen(
        [
            "kubectl",
            "-n",
            namespace,
            "port-forward",
            f"service/{service}",
            f"{local_port}:{service_port}",
        ]
    )

    def terminate_process():
        process.terminate()
        # Wait 10 seconds for the process to terminate.
        # This cleans up the zombie entry from the process table.
        # 10 seconds is a deliberately excessive amount of time to wait.
        process.wait(timeout=10)

    # Ensure clean-up in case of interrupt.
    atexit.register(terminate_process)
    # terminate_process is ok to execute multiple times.

    try:
        yield local_port
    finally:
        terminate_process()


def kubectl_patch(
    kind: str,
    name: str,
    namespace: str,
    patch: Dict[str, Any],
    patch_type: str = "strategic",
):
    """Wrapper for kubectl patch.

    Args:
        kind: Kind of the K8s resource (e.g. pod)
        name: Name of the K8s resource.
        namespace: Namespace of the K8s resource.
        patch: The patch to apply, as a dict.
        patch_type: json, merge, or strategic
    """
    with tempfile.NamedTemporaryFile("w") as patch_file:
        yaml.dump(patch, patch_file)
        patch_file.flush()
        subprocess.check_call(
            [
                "kubectl",
                "-n",
                f"{namespace}",
                "patch",
                f"{kind}",
                f"{name}",
                "--patch-file",
                f"{patch_file.name}",
                "--type",
                f"{patch_type}",
            ]
        )


def kubectl_delete(kind: str, name: str, namespace: str, wait: bool = True):
    """Wrapper for kubectl delete.

    Args:
        kind: Kind of the K8s resource (e.g. pod)
        name: Name of the K8s resource.
        namespace: Namespace of the K8s resource.
        wait: Whether to pass ``--wait=true`` so ``kubectl`` blocks until the
            resource is fully removed.
    """
    wait_str = "true" if wait else "false"
    subprocess.check_output(
        [
            "kubectl",
            "-n",
            f"{namespace}",
            "delete",
            f"{kind}",
            f"{name}",
            f"--wait={wait_str}",
        ]
    )