import sys
from collections import defaultdict

import pytest

import ray
from ray._common.network_utils import build_address
from ray._common.test_utils import (
    fetch_prometheus_metrics,
    wait_for_condition,
)
from ray._private.test_utils import run_string_as_driver_nonblocking

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


def raw_metrics(info):
    metrics_page = build_address("localhost", info["metrics_export_port"])
    res = fetch_prometheus_metrics([metrics_page])
    return res


def resources_by_state(info) -> dict:
    res = raw_metrics(info)
    if "ray_resources" in res:
        states = defaultdict(int)
        for sample in res["ray_resources"]:
            state = sample.labels["State"]
            name = sample.labels["Name"]
            states[(name, state)] += sample.value
            if states[(name, state)] == 0:
                del states[(name, state)]
        print("Resources by state: {}".format(states))
        return states
    else:
        return {}


def test_resources_metrics(shutdown_only):
    info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 3}, **METRIC_CONFIG)

    driver = """
import ray
import time

ray.init("auto")

@ray.remote(num_gpus=1, num_cpus=1, resources={"a": 1})
def f():
    time.sleep(999)

pg = ray.util.placement_group(bundles=[{"CPU": 1}])

ray.get([f.remote() for _ in range(2)])"""

    proc = run_string_as_driver_nonblocking(driver)

    def verify():
        state = resources_by_state(info)
        assert ("object_store_memory", "AVAILABLE") in state
        assert ("memory", "AVAILABLE") in state

        # Make sure node ip address & pg resources are not in state.
        node_ip_resource = None
        pg_resources = []

        for resource in ray.available_resources():
            if resource.startswith("node:"):
                node_ip_resource = resource
            elif "_group_" in resource:
                pg_resources.append(resource)

        # Verify node ip address resource not reported.
        assert (node_ip_resource, "AVAILABLE") not in state
        assert (node_ip_resource, "USED") not in state

        # Verify node pg resource not reported.
        for pg_resource in pg_resources:
            assert (pg_resource, "AVAILABLE") not in state
            assert (pg_resource, "USED") not in state

        assert state[("CPU", "AVAILABLE")] == 1.0
        # 2 tasks + 1 pg
        assert state[("CPU", "USED")] == 3.0
        assert ("GPU", "AVAILABLE") not in state  # since it is 0
        assert state[("GPU", "USED")] == 2.0

        # Verify custom resources
        assert state[("a", "AVAILABLE")] == 1.0
        assert state[("a", "USED")] == 2.0
        return True

    wait_for_condition(verify, timeout=20, retry_interval_ms=500)
    proc.kill()


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