import json
import logging
import os
import subprocess
import sys
from contextlib import contextmanager
from typing import Optional, Tuple

import pytest

import ray

logger = logging.getLogger(__name__)


@pytest.fixture
def shutdown_only():
    yield None
    # The code after the yield will run as teardown code.
    ray.shutdown()
    # Delete the cluster address just in case.
    ray._common.utils.reset_ray_address()


@contextmanager
def set_env_var(key: str, val: Optional[str] = None):
    old_val = os.environ.get(key, None)
    if val is not None:
        os.environ[key] = val
    elif key in os.environ:
        del os.environ[key]

    try:
        yield
    finally:
        if key in os.environ:
            del os.environ[key]
        if old_val is not None:
            os.environ[key] = old_val


@pytest.fixture
def ray_start_stop():
    subprocess.check_output(["ray", "start", "--head"])
    try:
        with set_env_var("RAY_ADDRESS", "http://127.0.0.1:8265"):
            yield
    finally:
        subprocess.check_output(["ray", "stop", "--force"])


@contextmanager
def ray_cluster_manager():
    """
    Used not as fixture in case we want to set RAY_ADDRESS first.
    """
    subprocess.check_output(["ray", "start", "--head"])
    try:
        yield
    finally:
        subprocess.check_output(["ray", "stop", "--force"])


def _run_cmd(cmd: str, should_fail=False) -> Tuple[str, str]:
    """Convenience wrapper for subprocess.run.

    We always run with shell=True to simulate the CLI.

    Asserts that the process succeeds/fails depending on should_fail.

    Returns (stdout, stderr).
    """
    print(f"Running command: '{cmd}'")
    p: subprocess.CompletedProcess = subprocess.run(
        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    if p.returncode == 0:
        print("Command succeeded.")
        if should_fail:
            raise RuntimeError(
                f"Expected command to fail, but got exit code: {p.returncode}."
            )
    else:
        print(f"Command failed with exit code: {p.returncode}.")
        if not should_fail:
            raise RuntimeError(
                f"Expected command to succeed, but got exit code: {p.returncode}."
            )

    return p.stdout.decode("utf-8"), p.stderr.decode("utf-8")


class TestJobSubmitHook:
    """Tests the RAY_JOB_SUBMIT_HOOK env var."""

    def test_hook(self, ray_start_stop):
        with set_env_var("RAY_JOB_SUBMIT_HOOK", "ray._private.test_utils.job_hook"):
            stdout, _ = _run_cmd("ray job submit -- echo hello")
            assert "hook intercepted: echo hello" in stdout


class TestRayJobHeaders:
    """
    Integration version of job CLI test that ensures interaction with the
    following components are working as expected:
    1) Ray client: use of RAY_JOB_HEADERS and ray.init() in job_head.py
    2) Ray dashboard: `ray start --head`
    """

    def test_empty_ray_job_headers(self, ray_start_stop):
        with set_env_var("RAY_JOB_HEADERS", None):
            stdout, _ = _run_cmd("ray job submit -- echo hello")
            assert "hello" in stdout
            assert "succeeded" in stdout

    @pytest.mark.parametrize("ray_job_headers", ['{"key": "value"}'])
    def test_ray_job_headers(self, ray_start_stop, ray_job_headers: str):
        with set_env_var("RAY_JOB_HEADERS", ray_job_headers):
            _run_cmd("ray job submit -- echo hello", should_fail=False)

    @pytest.mark.parametrize("ray_job_headers", ["{key value}"])
    def test_ray_incorrectly_formatted_job_headers(
        self, ray_start_stop, ray_job_headers: str
    ):
        with set_env_var("RAY_JOB_HEADERS", ray_job_headers):
            _run_cmd("ray job submit -- echo hello", should_fail=True)


class TestRayAddress:
    """
    Integration version of job CLI test that ensures interaction with the
    following components are working as expected:

    1) Ray client: use of RAY_ADDRESS and ray.init() in job_head.py
    2) Ray dashboard: `ray start --head`
    """

    def test_empty_ray_address(self, ray_start_stop):
        with set_env_var("RAY_ADDRESS", None):
            stdout, _ = _run_cmd("ray job submit -- echo hello")
            assert "hello" in stdout
            assert "succeeded" in stdout

    @pytest.mark.parametrize(
        "ray_api_server_address,should_fail",
        [
            ("http://127.0.0.1:8265", False),  # correct API server
            ("127.0.0.1:8265", True),  # wrong format without http
            ("http://127.0.0.1:9999", True),  # wrong port
        ],
    )
    def test_ray_api_server_address(
        self,
        ray_start_stop,
        ray_api_server_address: str,
        should_fail: bool,
    ):
        # Set a `RAY_ADDRESS` that would not work with the `ray job submit` CLI because it uses the `ray://` prefix.
        # This verifies that the `RAY_API_SERVER_ADDRESS` env var takes precedence.
        with set_env_var("RAY_ADDRESS", "ray://127.0.0.1:8265"):
            with set_env_var("RAY_API_SERVER_ADDRESS", ray_api_server_address):
                _run_cmd("ray job submit -- echo hello", should_fail=should_fail)

    @pytest.mark.parametrize(
        "ray_client_address,should_fail",
        [
            ("127.0.0.1:8265", True),
            ("ray://127.0.0.1:8265", True),
            ("http://127.0.0.1:8265", False),
        ],
    )
    def test_ray_client_address(
        self, ray_start_stop, ray_client_address: str, should_fail: bool
    ):
        with set_env_var("RAY_ADDRESS", ray_client_address):
            _run_cmd("ray job submit -- echo hello", should_fail=should_fail)

    def test_valid_http_ray_address(self, ray_start_stop):
        stdout, _ = _run_cmd("ray job submit -- echo hello")
        assert "hello" in stdout
        assert "succeeded" in stdout


class TestJobSubmit:
    def test_basic_submit(self, ray_start_stop):
        """Should tail logs and wait for process to exit."""
        cmd = "sleep 1 && echo hello && sleep 1 && echo hello"
        stdout, _ = _run_cmd(f"ray job submit -- bash -c '{cmd}'")

        # 'hello' should appear four times: twice when we print the entrypoint, then
        # two more times in the logs from the `echo`.
        assert stdout.count("hello") == 4
        assert "succeeded" in stdout

    def test_submit_no_wait(self, ray_start_stop):
        """Should exit immediately w/o printing logs."""
        cmd = "echo hello && sleep 1000"
        stdout, _ = _run_cmd(f"ray job submit --no-wait -- bash -c '{cmd}'")
        assert "hello" not in stdout
        assert "Tailing logs until the job exits" not in stdout

    def test_submit_with_logs_instant_job(self, ray_start_stop):
        """Should exit immediately and print logs even if job returns instantly."""
        cmd = "echo hello"
        stdout, _ = _run_cmd(f"ray job submit -- bash -c '{cmd}'")

        # 'hello' should appear twice: once when we print the entrypoint, then
        # again from the `echo`.
        assert stdout.count("hello") == 2

    def test_multiple_ray_init(self, ray_start_stop):
        cmd = (
            "python -c 'import ray; ray.init(); ray.shutdown(); "
            "ray.init(); ray.shutdown();'"
        )
        stdout, _ = _run_cmd(f"ray job submit -- {cmd}")
        assert "succeeded" in stdout

    def test_metadata(self, ray_start_stop):
        cmd = "echo hello"
        stdout, _ = _run_cmd(
            f'ray job submit --metadata-json=\'{{"key": "value"}}\' -- {cmd}'
        )
        assert "hello" in stdout
        assert "succeeded" in stdout

    def test_job_failed(self, ray_start_stop):
        cmd = "python -c 'import ray; ray.init(); assert 1 == 2;'"
        _run_cmd(f"ray job submit -- {cmd}", should_fail=True)


class TestRuntimeEnv:
    def test_bad_runtime_env(self, ray_start_stop):
        """Should fail with helpful error if runtime env setup fails."""
        stdout, _ = _run_cmd(
            'ray job submit --runtime-env-json=\'{"pip": '
            '["does-not-exist"]}\' -- echo hi',
            should_fail=True,
        )
        assert "Tailing logs until the job exits" in stdout
        assert "runtime_env setup failed" in stdout
        assert "No matching distribution found for does-not-exist" in stdout


class TestJobStop:
    def test_basic_stop(self, ray_start_stop):
        """Should wait until the job is stopped."""
        cmd = "sleep 1000"
        job_id = "test_basic_stop"
        _run_cmd(f"ray job submit --no-wait --job-id={job_id} -- {cmd}")

        stdout, _ = _run_cmd(f"ray job stop {job_id}")
        assert "Waiting for job" in stdout
        assert f"Job '{job_id}' was stopped" in stdout

    def test_stop_no_wait(self, ray_start_stop):
        """Should not wait until the job is stopped."""
        cmd = "echo hello && sleep 1000"
        job_id = "test_stop_no_wait"
        _run_cmd(f"ray job submit --no-wait --job-id={job_id} -- bash -c '{cmd}'")

        stdout, _ = _run_cmd(f"ray job stop --no-wait {job_id}")
        assert "Waiting for job" not in stdout
        assert f"Job '{job_id}' was stopped" not in stdout


class TestJobList:
    def test_empty(self, ray_start_stop):
        stdout, _ = _run_cmd("ray job list")
        assert "[]" in stdout

    def test_list(self, ray_start_stop):
        _run_cmd("ray job submit --job-id='hello_id' -- echo hello")

        runtime_env = {"env_vars": {"TEST": "123"}}
        _run_cmd(
            "ray job submit --job-id='hi_id' "
            f"--runtime-env-json='{json.dumps(runtime_env)}' -- echo hi"
        )
        stdout, _ = _run_cmd("ray job list")
        assert "123" in stdout
        assert "hello_id" in stdout
        assert "hi_id" in stdout


class TestJobDelete:
    def test_basic_delete(self, ray_start_stop):
        cmd = "sleep 1000"
        job_id = "test_basic_delete"
        _run_cmd(f"ray job submit --no-wait --submission-id={job_id} -- {cmd}")

        # Job shouldn't be able to be deleted because it is not in a terminal state.
        stdout, stderr = _run_cmd(f"ray job delete {job_id}", should_fail=True)
        assert "it is in a non-terminal state" in stderr

        # Submit a job that finishes quickly.
        cmd = "echo hello"
        job_id = "test_basic_delete_quick"
        _run_cmd(f"ray job submit --submission-id={job_id} -- bash -c '{cmd}'")

        # Job should be able to be deleted because it is finished.
        stdout, _ = _run_cmd(f"ray job delete {job_id}")
        assert f"Job '{job_id}' deleted successfully" in stdout


class TestJobStatus:
    # `ray job status` should exit with 0 if the job exists and non-zero if it doesn't.
    # This is the contract between Ray and KubRay v1.3.0.
    def test_status_job_exists(self, ray_start_stop):
        cmd = "echo hello"
        job_id = "test_job_id"
        _run_cmd(
            f"ray job submit --submission-id={job_id} -- bash -c '{cmd}'",
            should_fail=False,
        )
        _run_cmd(f"ray job status {job_id}", should_fail=False)

    def test_status_job_does_not_exist(self, ray_start_stop):
        job_id = "test_job_id"
        _run_cmd(f"ray job status {job_id}", should_fail=True)


def test_quote_escaping(ray_start_stop):
    cmd = "echo \"hello 'world'\""
    job_id = "test_quote_escaping"
    stdout, _ = _run_cmd(
        f"ray job submit --job-id={job_id} -- {cmd}",
    )
    assert "hello 'world'" in stdout


def test_resources(shutdown_only):
    ray.init(num_cpus=1, num_gpus=1, resources={"Custom": 1}, _memory=4)

    # Check the case of too many resources.
    for id, arg in [
        ("entrypoint_num_cpus", "--entrypoint-num-cpus=2"),
        ("entrypoint_num_gpus", "--entrypoint-num-gpus=2"),
        ("entrypoint_memory", "--entrypoint-memory=5"),
        ("entrypoint_resources", "--entrypoint-resources='{\"Custom\": 2}'"),
    ]:
        _run_cmd(f"ray job submit --submission-id={id} --no-wait {arg} -- echo hi")
        stdout, _ = _run_cmd(f"ray job status {id}")
        assert "waiting for resources" in stdout

    # Check the case of sufficient resources.
    stdout, _ = _run_cmd(
        "ray job submit --entrypoint-num-cpus=1 "
        "--entrypoint-num-gpus=1 --entrypoint-memory=4 --entrypoint-resources='{"
        '"Custom": 1}\' -- echo hello',
    )
    assert "hello" in stdout


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", __file__]))