import copy
import multiprocessing
import sys
from collections import defaultdict

import pytest

import ray
from ray._common.test_utils import (
    PrometheusTimeseries,
    run_string_as_driver,
    wait_for_condition,
)
from ray._private.metrics_agent import RAY_WORKER_TIMEOUT_S
from ray._private.test_utils import (
    raw_metric_timeseries,
    run_string_as_driver_nonblocking,
    wait_for_assertion,
    wait_for_dashboard_agent_available,
)

METRIC_CONFIG = {
    "_system_config": {
        "metrics_report_interval_ms": 100,
    }
}

SLOW_METRIC_CONFIG = {
    "_system_config": {
        "metrics_report_interval_ms": 3000,
    }
}


def tasks_by_state(info, timeseries: PrometheusTimeseries, flush: bool = False) -> dict:
    if flush:
        timeseries.flush()
    return tasks_breakdown(info, lambda s: s.labels["State"], timeseries)


def tasks_by_name_and_state(info, timeseries: PrometheusTimeseries) -> dict:
    return tasks_breakdown(
        info, lambda s: (s.labels["Name"], s.labels["State"]), timeseries
    )


def tasks_by_all(info, timeseries: PrometheusTimeseries) -> dict:
    return tasks_breakdown(
        info,
        lambda s: (s.labels["Name"], s.labels["State"], s.labels["IsRetry"]),
        timeseries,
    )


def tasks_breakdown(info, key_fn, timeseries: PrometheusTimeseries) -> dict:
    res = raw_metric_timeseries(info, timeseries)
    if "ray_tasks" in res:
        breakdown = defaultdict(int)
        for sample in res["ray_tasks"]:
            key = key_fn(sample)
            breakdown[key] += sample.value
            if breakdown[key] == 0:
                del breakdown[key]
        print("Task label breakdown: {}".format(breakdown))
        return breakdown
    else:
        return {}


# TODO(ekl) in all these tests, we use run_string_as_driver_nonblocking to work around
# stats reporting issues if Ray is repeatedly restarted in unit tests.
def test_task_basic(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)

    driver = """
import time

import ray
from ray._common.test_utils import wait_for_condition

ray.init("auto")

@ray.remote
def a():
    time.sleep(999)

@ray.remote
def b():
    time.sleep(999)

@ray.remote(num_cpus=3)
def c():
    time.sleep(999)

refs = [a.remote(), b.remote()]
wait_for_condition(
    lambda: ray.available_resources().get("CPU", 0) == 0,
)

ray.get(refs + [c.remote() for _ in range(8)])
"""
    proc = run_string_as_driver_nonblocking(driver)
    timeseries = PrometheusTimeseries()
    expected = {
        "RUNNING": 2.0,
        "PENDING_NODE_ASSIGNMENT": 8.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )
    assert tasks_by_name_and_state(info, timeseries) == {
        ("a", "RUNNING"): 1.0,
        ("b", "RUNNING"): 1.0,
        ("c", "PENDING_NODE_ASSIGNMENT"): 8.0,
    }
    proc.kill()


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_task_custom_name_metrics(shutdown_only):
    """Verify that custom task names set via .options(name=...) are used in metrics.

    This tests that RUNNING tasks use the custom name consistently with
    FINISHED/FAILED tasks. Previously there was a bug where RUNNING metrics used
    the function name (FunctionDescriptor->CallString()) but FINISHED/FAILED used
    the custom name (TaskSpec::GetName()).
    """
    info = ray.init(num_cpus=2, **METRIC_CONFIG)

    driver = """
import ray
import time

ray.init("auto")

@ray.remote
def my_function():
    time.sleep(999)

# Submit tasks with custom names
a = [my_function.options(name="custom_task_name").remote() for _ in range(4)]
ray.get(a)
"""
    proc = run_string_as_driver_nonblocking(driver)
    timeseries = PrometheusTimeseries()

    # Verify that RUNNING tasks use the custom name, not the function name.
    # With 2 CPUs, 2 tasks should be running and 2 should be pending.
    expected = {
        ("custom_task_name", "RUNNING"): 2.0,
        ("custom_task_name", "PENDING_NODE_ASSIGNMENT"): 2.0,
    }
    wait_for_condition(
        lambda: tasks_by_name_and_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )

    # Verify the original function name is NOT used in metrics
    breakdown = tasks_by_name_and_state(info, timeseries)
    assert (
        "my_function",
        "RUNNING",
    ) not in breakdown, "RUNNING tasks should use custom name, not function name"
    assert (
        "my_function",
        "PENDING_NODE_ASSIGNMENT",
    ) not in breakdown, "PENDING tasks should use custom name, not function name"

    proc.kill()


def test_task_job_ids(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote(num_cpus=0)
def {func_name}():
    time.sleep(999)
a = [{func_name}.remote() for _ in range(1)]
ray.get(a)
"""
    # We make sure the task name is unique for each job so that we can distinguish the metric sample by task name across jobs.
    procs = [
        run_string_as_driver_nonblocking(driver.format(func_name=f"f_{i}"))
        for i in range(3)
    ]

    expected = {
        "RUNNING": 3.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )

    # Check we have three jobs reporting "RUNNING".
    metrics = raw_metric_timeseries(info, timeseries)
    jobs_at_state = defaultdict(set)
    for sample in metrics["ray_tasks"]:
        jobs_at_state[sample.labels["State"]].add(sample.labels["JobId"])
    print("Jobs at state: {}".format(jobs_at_state))
    assert len(jobs_at_state["RUNNING"]) == 3, jobs_at_state

    for proc in procs:
        proc.kill()


def test_task_nested(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote(num_cpus=0)
def wrapper():
    @ray.remote
    def a():
        time.sleep(999)

    @ray.remote
    def b():
        time.sleep(999)

    @ray.remote
    def c():
        time.sleep(999)

    ray.get([a.remote(), b.remote()] + [c.remote() for _ in range(8)])

w = wrapper.remote()
ray.get(w)
"""
    proc = run_string_as_driver_nonblocking(driver)

    expected = {
        "RUNNING": 2.0,
        "RUNNING_IN_RAY_GET": 1.0,
        "PENDING_NODE_ASSIGNMENT": 8.0,
    }

    def check_task_state():
        assert tasks_by_state(info, timeseries) == expected

    wait_for_assertion(
        check_task_state,
        timeout=30,
        retry_interval_ms=2000,
    )
    assert tasks_by_name_and_state(info, timeseries) == {
        ("wrapper", "RUNNING_IN_RAY_GET"): 1.0,
        ("a", "RUNNING"): 1.0,
        ("b", "RUNNING"): 1.0,
        ("c", "PENDING_NODE_ASSIGNMENT"): 8.0,
    }
    proc.kill()


def test_task_nested_wait(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote(num_cpus=0)
def wrapper():
    @ray.remote
    def a():
        time.sleep(999)

    @ray.remote
    def b():
        time.sleep(999)

    @ray.remote
    def c():
        time.sleep(999)

    ray.wait([a.remote(), b.remote()] + [c.remote() for _ in range(8)])

w = wrapper.remote()
ray.get(w)
"""
    proc = run_string_as_driver_nonblocking(driver)

    expected = {
        "RUNNING": 2.0,
        "RUNNING_IN_RAY_WAIT": 1.0,
        "PENDING_NODE_ASSIGNMENT": 8.0,
    }

    def check_task_state():
        assert tasks_by_state(info, timeseries) == expected

    wait_for_assertion(
        check_task_state,
        timeout=30,
        retry_interval_ms=2000,
    )
    assert tasks_by_name_and_state(info, timeseries) == {
        ("wrapper", "RUNNING_IN_RAY_WAIT"): 1.0,
        ("a", "RUNNING"): 1.0,
        ("b", "RUNNING"): 1.0,
        ("c", "PENDING_NODE_ASSIGNMENT"): 8.0,
    }
    proc.kill()


def driver_for_test_task_fetch_args(head_info):
    ray.init("auto")
    timeseries = PrometheusTimeseries()

    @ray.remote(resources={"worker": 1})
    def task1():
        return [1] * 1024 * 1024

    @ray.remote(resources={"head": 1})
    def task2(obj):
        pass

    o1 = task1.remote()
    o2 = task2.remote(o1)

    wait_for_condition(
        lambda: tasks_by_state(head_info, timeseries).get("PENDING_ARGS_FETCH", 0.0)
        == 1.0
    )

    ray.cancel(o2)

    wait_for_condition(
        lambda: tasks_by_state(head_info, timeseries).get("PENDING_ARGS_FETCH", 0.0)
        == 0.0
    )


def test_task_fetch_args(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(
        resources={"head": 1},
        _system_config={
            "metrics_report_interval_ms": 100,
            "testing_asio_delay_us": "ObjectManagerService.grpc_server.Pull=5000000000:5000000000",  # noqa: E501
        },
    )
    head_info = ray.init(address=cluster.address)
    cluster.add_node(resources={"worker": 1})
    cluster.wait_for_nodes()

    multiprocessing.set_start_method("spawn")
    p = multiprocessing.Process(
        target=driver_for_test_task_fetch_args, args=(head_info,)
    )
    p.start()
    p.join()
    assert p.exitcode == 0


def test_task_wait_on_deps(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote
def f():
    time.sleep(999)

@ray.remote
def g(x):
    time.sleep(999)

x = f.remote()
a = [g.remote(x) for _ in range(5)]
ray.get(a)
"""
    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        "RUNNING": 1.0,
        "PENDING_ARGS_AVAIL": 5.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )
    assert tasks_by_name_and_state(info, timeseries) == {
        ("f", "RUNNING"): 1.0,
        ("g", "PENDING_ARGS_AVAIL"): 5.0,
    }
    proc.kill()


def test_actor_tasks_queued(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote
class F:
    def f(self):
        time.sleep(999)

    def g(self):
        pass

a = F.remote()
[a.g.remote() for _ in range(10)]
[a.f.remote() for _ in range(1)]  # Further tasks should be blocked on this one.
z = [a.g.remote() for _ in range(9)]
ray.get(z)
"""
    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        ("F.__init__", "FINISHED"): 1.0,
        ("F.g", "FINISHED"): 10.0,
        ("F.f", "RUNNING"): 1.0,
        ("F.g", "SUBMITTED_TO_WORKER"): 9.0,
    }
    wait_for_condition(
        lambda: tasks_by_name_and_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


def test_task_finish(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote
def f():
    return "ok"

@ray.remote
def g():
    assert False

f.remote()
g.remote()
time.sleep(999)
"""

    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        "FAILED": 1.0,
        "FINISHED": 1.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )
    assert tasks_by_name_and_state(info, timeseries) == {
        ("g", "FAILED"): 1.0,
        ("f", "FINISHED"): 1.0,
    }
    proc.kill()


def test_task_retry(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote
def sleep():
    time.sleep(999)

@ray.remote
class Phaser:
    def __init__(self):
        self.i = 0

    def inc(self):
        self.i += 1
        if self.i < 3:
            raise ValueError("First two tries will fail")

phaser = Phaser.remote()

@ray.remote(retry_exceptions=True, max_retries=3)
def f():
    ray.get(phaser.inc.remote())
    ray.get(sleep.remote())

f.remote()
time.sleep(999)
"""

    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        ("sleep", "RUNNING", "0"): 1.0,
        ("f", "FAILED", "0"): 1.0,
        ("f", "FAILED", "1"): 1.0,
        ("f", "RUNNING_IN_RAY_GET", "1"): 1.0,
        ("Phaser.__init__", "FINISHED", "0"): 1.0,
        ("Phaser.inc", "FINISHED", "0"): 1.0,
        ("Phaser.inc", "FAILED", "0"): 2.0,
    }
    wait_for_condition(
        lambda: expected.items() <= tasks_by_all(info, timeseries).items(),
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows. Timing out.")
def test_actor_task_retry(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import os
import time

ray.init("auto")

@ray.remote
class Phaser:
    def __init__(self):
        self.i = 0

    def inc(self):
        self.i += 1
        if self.i < 3:
            raise ValueError("First two tries will fail")

phaser = Phaser.remote()

@ray.remote(max_restarts=10, max_task_retries=10)
class F:
    def f(self):
        try:
            ray.get(phaser.inc.remote())
        except Exception:
            print("RESTART")
            os._exit(1)

f = F.remote()
ray.get(f.f.remote())
time.sleep(999)
"""

    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        ("F.__init__", "FINISHED", "0"): 1.0,
        ("F.f", "FAILED", "0"): 1.0,
        ("F.f", "FAILED", "1"): 1.0,
        ("F.f", "FINISHED", "1"): 1.0,
        ("Phaser.__init__", "FINISHED", "0"): 1.0,
        ("Phaser.inc", "FINISHED", "0"): 1.0,
    }
    wait_for_condition(
        lambda: expected.items() <= tasks_by_all(info, timeseries).items(),
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_task_failure(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time
import os

ray.init("auto")

@ray.remote(max_retries=0)
def f():
    print("RUNNING FAILING TASK")
    os._exit(1)

@ray.remote
def g():
    assert False

f.remote()
g.remote()
time.sleep(999)
"""

    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        "FAILED": 2.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


def test_concurrent_actor_tasks(shutdown_only):
    info = ray.init(num_cpus=2, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import asyncio

ray.init("auto")

@ray.remote(max_concurrency=30)
class A:
    async def f(self):
        await asyncio.sleep(300)

a = A.remote()
ray.get([a.f.remote() for _ in range(40)])
"""

    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        "RUNNING": 30.0,
        "SUBMITTED_TO_WORKER": 10.0,
        "FINISHED": 1.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_export_now(shutdown_only, ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(
        **SLOW_METRIC_CONFIG,
        num_cpus=2,
    )
    wait_for_dashboard_agent_available(cluster)
    info = ray.init(address=cluster.address)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time

ray.init("auto")

@ray.remote
def {func_name}():
    pass
a = [{func_name}.remote() for _ in range(10)]
ray.get(a)
"""

    # If force export at process death is broken, we won't see the recently completed
    # tasks from the drivers. We also make sure the task name is unique for each job so
    # that we can distinguish the metric sample by task name across jobs.
    for i in range(10):
        print("Run job", i)
        run_string_as_driver(driver.format(func_name=f"f_{i}"))
        tasks_by_state(info, timeseries)

    expected = {
        "FINISHED": 100.0,
    }
    wait_for_condition(
        lambda: tasks_by_state(info, timeseries) == expected,
        timeout=20,
        retry_interval_ms=500,
    )


@pytest.mark.skipif(sys.platform == "darwin", reason="Flaky on macos")
def test_pull_manager_stats(shutdown_only):
    info = ray.init(num_cpus=2, object_store_memory=100_000_000, **METRIC_CONFIG)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import time
import numpy as np

ray.init("auto")

# Spill a lot of 10MiB objects. The object store is 100MiB, so pull manager will
# only be able to pull ~9 total into memory at once, including running tasks.
buf = []
for _ in range(100):
    buf.append(ray.put(np.ones(10 * 1024 * 1024, dtype=np.uint8)))

@ray.remote
def a(x):
    time.sleep(999)

@ray.remote
def b(x):
    time.sleep(999)

@ray.remote
def c(x):
    time.sleep(999)

ray.get([a.remote(buf[0]), b.remote(buf[1])] + [c.remote(x) for x in buf[2:]])
"""

    proc = run_string_as_driver_nonblocking(driver)

    # This test is non-deterministic since pull bundles can sometimes end up fallback
    # allocated. This leads to slightly more objects pulled than you'd expect.
    def close_to_expected(stats):
        # A scheduled task can momentarily sit in SUBMITTED_TO_WORKER (lease granted,
        # waiting on the worker to fetch its spilled arg) before it reports RUNNING.
        # Under the object store pressure this test creates, that transition can be
        # slow, so count SUBMITTED_TO_WORKER as running to avoid flakiness.
        running = stats.get("RUNNING", 0) + stats.get("SUBMITTED_TO_WORKER", 0)
        assert running == 2, stats
        assert 7 <= stats["PENDING_NODE_ASSIGNMENT"] <= 17, stats
        assert 81 <= stats["PENDING_OBJ_STORE_MEM_AVAIL"] <= 91, stats
        assert set(stats.keys()).issubset(
            {
                "RUNNING",
                "SUBMITTED_TO_WORKER",
                "PENDING_NODE_ASSIGNMENT",
                "PENDING_OBJ_STORE_MEM_AVAIL",
            }
        ), stats
        assert sum(stats.values()) == 100, stats
        return True

    wait_for_condition(
        lambda: close_to_expected(tasks_by_state(info, timeseries)),
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_stale_view_cleanup_when_job_exits(monkeypatch, shutdown_only):
    timeseries = PrometheusTimeseries()
    with monkeypatch.context() as m:
        m.setenv(RAY_WORKER_TIMEOUT_S, 5)
        info = ray.init(num_cpus=2, **METRIC_CONFIG)
        print(info)

        driver = """
import ray
import time
import numpy as np

ray.init("auto")

@ray.remote
def g():
    time.sleep(999)

ray.get(g.remote())
    """

        proc = run_string_as_driver_nonblocking(driver)
        expected = {
            "RUNNING": 1.0,
        }
        wait_for_condition(
            lambda: tasks_by_state(info, timeseries) == expected,
            timeout=20,
            retry_interval_ms=500,
        )

        proc.kill()
        print("Killing a driver.")
        expected = {}
        wait_for_condition(
            lambda: tasks_by_state(info, timeseries, flush=True) == expected,
            timeout=20,
            retry_interval_ms=500,
        )


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows. Timing out.")
def test_metrics_batch(shutdown_only):
    """Verify metrics_report_batch_size works correctly without data loss."""
    config_copy = copy.deepcopy(METRIC_CONFIG)
    config_copy["_system_config"].update({"metrics_report_batch_size": 1})
    info = ray.init(num_cpus=2, **config_copy)
    timeseries = PrometheusTimeseries()
    driver = """
import ray
import os
import time

ray.init("auto")

@ray.remote
class Phaser:
    def __init__(self):
        self.i = 0

    def inc(self):
        self.i += 1
        if self.i < 3:
            raise ValueError("First two tries will fail")

phaser = Phaser.remote()

@ray.remote(max_restarts=10, max_task_retries=10)
class F:
    def f(self):
        try:
            ray.get(phaser.inc.remote())
        except Exception:
            print("RESTART")
            os._exit(1)

f = F.remote()
ray.get(f.f.remote())
time.sleep(999)
"""

    proc = run_string_as_driver_nonblocking(driver)
    expected = {
        ("F.__init__", "FINISHED", "0"): 1.0,
        ("F.f", "FAILED", "0"): 1.0,
        ("F.f", "FAILED", "1"): 1.0,
        ("F.f", "FINISHED", "1"): 1.0,
        ("Phaser.__init__", "FINISHED", "0"): 1.0,
        ("Phaser.inc", "FINISHED", "0"): 1.0,
    }
    wait_for_condition(
        lambda: expected.items() <= tasks_by_all(info, timeseries).items(),
        timeout=20,
        retry_interval_ms=500,
    )
    proc.kill()


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