import random
import sys
import time
from typing import List

import pytest
from opencensus.metrics.export.metric_descriptor import MetricDescriptorType
from opencensus.metrics.export.value import ValueDouble
from opencensus.stats import execution_context
from opencensus.stats.aggregation_data import (
    CountAggregationData,
    DistributionAggregationData,
    LastValueAggregationData,
    SumAggregationData,
)
from opencensus.stats.stats_recorder import StatsRecorder
from opencensus.stats.view_manager import ViewManager
from prometheus_client.core import REGISTRY

import ray._private.prometheus_exporter as prometheus_exporter
from ray._common.test_utils import (
    fetch_prometheus_metrics,
    fetch_raw_prometheus,
    wait_for_condition,
)
from ray._private.metrics_agent import (
    RAY_WORKER_TIMEOUT_S,
    Gauge,
    MetricsAgent,
    OpenCensusProxyCollector,
    OpencensusProxyMetric,
    Record,
)
from ray._private.telemetry.metric_cardinality import WORKER_ID_TAG_KEY
from ray._raylet import WorkerID
from ray.core.generated.metrics_pb2 import (
    LabelKey,
    LabelValue,
    Metric,
    MetricDescriptor,
    Point,
    TimeSeries,
)


def raw_metrics(export_port):
    metrics_page = "localhost:{}".format(export_port)
    res = fetch_prometheus_metrics([metrics_page])
    return res


def get_metric(metric_name, export_port):
    res = raw_metrics(export_port)
    for name, samples in res.items():
        if name == metric_name:
            return name, samples

    return None


def get_prom_metric_name(namespace, metric_name):
    return f"{namespace}_{metric_name}"


def generate_timeseries(label_values: List[str], points: List[float]):
    return TimeSeries(
        label_values=[LabelValue(value=val) for val in label_values],
        points=[Point(double_value=val) for val in points],
    )


def generate_protobuf_metric(
    name: str,
    desc: str,
    unit: str,
    type: MetricDescriptorType,
    label_keys: List[str] = None,
    timeseries: List[TimeSeries] = None,
):
    if not label_keys:
        label_keys = []
    if not timeseries:
        timeseries = []

    return Metric(
        metric_descriptor=MetricDescriptor(
            name=name,
            description=desc,
            unit=unit,
            type=type,
            label_keys=[LabelKey(key="a"), LabelKey(key="b")],
        ),
        timeseries=timeseries,
    )


@pytest.fixture
def get_agent(request, monkeypatch):
    with monkeypatch.context() as m:
        if hasattr(request, "param"):
            delay = request.param
        else:
            delay = 0

        m.setenv(RAY_WORKER_TIMEOUT_S, delay)
        agent_port = random.randint(10000, 65535)
        stats_recorder = StatsRecorder()
        view_manager = ViewManager()
        stats_exporter = prometheus_exporter.new_stats_exporter(
            prometheus_exporter.Options(
                namespace="test",
                port=agent_port,
                address="127.0.0.1",
            )
        )
        agent = MetricsAgent(view_manager, stats_recorder, stats_exporter)
        REGISTRY.register(agent.proxy_exporter_collector)
        yield agent, agent_port
        REGISTRY.unregister(agent.stats_exporter.collector)
        REGISTRY.unregister(agent.proxy_exporter_collector)
        execution_context.set_measure_to_view_map({})


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_agent_record_and_export(get_agent):
    namespace = "test"
    agent, agent_port = get_agent

    # Record a new gauge.
    metric_name = "test"
    test_gauge = Gauge(metric_name, "desc", "unit", ["tag"])
    record_a = Record(
        gauge=test_gauge,
        value=3,
        tags={"tag": "a"},
    )
    agent.record_and_export([record_a])
    name, samples = get_metric(get_prom_metric_name(namespace, metric_name), agent_port)
    assert name == get_prom_metric_name(namespace, metric_name)
    assert len(samples) == 1
    assert samples[0].value == 3
    assert samples[0].labels == {"tag": "a"}

    # Record the same gauge.
    record_b = Record(
        gauge=test_gauge,
        value=4,
        tags={"tag": "a"},
    )
    record_c = Record(
        gauge=test_gauge,
        value=4,
        tags={"tag": "a"},
    )
    agent.record_and_export([record_b, record_c])
    name, samples = get_metric(get_prom_metric_name(namespace, metric_name), agent_port)
    assert name == get_prom_metric_name(namespace, metric_name)
    assert len(samples) == 1
    assert samples[0].value == 4
    assert samples[0].labels == {"tag": "a"}

    # Record the same gauge with different ag.
    record_d = Record(
        gauge=test_gauge,
        value=6,
        tags={"tag": "aa"},
    )
    agent.record_and_export(
        [
            record_d,
        ]
    )
    name, samples = get_metric(get_prom_metric_name(namespace, metric_name), agent_port)
    assert name == get_prom_metric_name(namespace, metric_name)
    assert len(samples) == 2
    assert samples[0].value == 4
    assert samples[0].labels == {"tag": "a"}
    assert samples[1].value == 6
    assert samples[1].labels == {"tag": "aa"}

    # Record more than 1 gauge.
    metric_name_2 = "test2"
    test_gauge_2 = Gauge(metric_name_2, "desc", "unit", ["tag"])
    record_e = Record(
        gauge=test_gauge_2,
        value=1,
        tags={"tag": "b"},
    )
    agent.record_and_export([record_e])
    name, samples = get_metric(
        get_prom_metric_name(namespace, metric_name_2), agent_port
    )

    assert name == get_prom_metric_name(namespace, metric_name_2)
    assert samples[0].value == 1
    assert samples[0].labels == {"tag": "b"}

    # Make sure the previous record is still there.
    name, samples = get_metric(get_prom_metric_name(namespace, metric_name), agent_port)
    assert name == get_prom_metric_name(namespace, metric_name)
    assert len(samples) == 2
    assert samples[0].value == 4
    assert samples[0].labels == {"tag": "a"}
    assert samples[1].value == 6
    assert samples[1].labels == {"tag": "aa"}


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_agent_record_and_export_failed_records_dont_block_other_records(
    get_agent,
    capsys,
):
    namespace = "test"
    agent, agent_port = get_agent

    metric_name = "test"
    test_gauge = Gauge(metric_name, "desc", "unit", ["tag"])
    record_a = Record(
        gauge=test_gauge,
        value=1,
        tags={"tag": "a"},
    )
    record_b = Record(
        gauge=test_gauge,
        value=1,
        # this tag is much too long (>255 characters), so recording this metric will fail
        tags={"tag": "b" * 1000},
    )
    record_c = Record(
        gauge=test_gauge,
        value=1,
        tags={"tag": "c"},
    )
    agent.record_and_export([record_a, record_b, record_c])

    name, samples = get_metric(get_prom_metric_name(namespace, metric_name), agent_port)
    assert name == get_prom_metric_name(namespace, metric_name)

    # a and c should be recorded, b's failure should be ignored
    assert {sample.labels["tag"] for sample in samples} == {"a", "c"}


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_agent_proxy_record_and_export_basic(get_agent):
    """Test the case the metrics are exported without worker_id."""
    namespace = "test"
    agent, agent_port = get_agent

    # Test the basic case.
    m = generate_protobuf_metric(
        "test",
        "desc",
        "",
        MetricDescriptorType.GAUGE_DOUBLE,
        label_keys=["a", "b"],
        timeseries=[],
    )
    m.timeseries.append(generate_timeseries(["a", "b"], [1, 2, 3]))
    agent.proxy_export_metrics([m])
    name, samples = get_metric(f"{namespace}_test", agent_port)
    assert name == f"{namespace}_test"
    assert len(samples) == 1
    assert samples[0].labels == {"a": "a", "b": "b"}
    assert samples[0].value == 3

    # Test new metric has proxyed.
    m = generate_protobuf_metric(
        "test",
        "desc",
        "",
        MetricDescriptorType.GAUGE_DOUBLE,
        label_keys=["a", "b"],
        timeseries=[],
    )
    m.timeseries.append(generate_timeseries(["a", "b"], [4]))
    agent.proxy_export_metrics([m])
    name, samples = get_metric(f"{namespace}_test", agent_port)
    assert name == f"{namespace}_test"
    assert len(samples) == 1
    assert samples[0].labels == {"a": "a", "b": "b"}
    assert samples[0].value == 4

    # Test new metric with different tag is reported.
    m = generate_protobuf_metric(
        "test",
        "desc",
        "",
        MetricDescriptorType.GAUGE_DOUBLE,
        label_keys=["a", "b"],
        timeseries=[],
    )
    m.timeseries.append(generate_timeseries(["a", "c"], [5]))
    agent.proxy_export_metrics([m])
    name, samples = get_metric(f"{namespace}_test", agent_port)
    assert name == f"{namespace}_test"
    assert len(samples) == 2
    assert samples[0].labels == {"a": "a", "b": "b"}
    assert samples[0].value == 4
    # Newly added metric has different tags and values.
    assert samples[1].labels == {"a": "a", "b": "c"}
    assert samples[1].value == 5


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_agent_proxy_record_and_export_from_workers(get_agent):
    """
    Test the basic worker death case.
    """
    namespace = "test"
    agent, agent_port = get_agent
    worker_id = WorkerID.from_random()

    m = generate_protobuf_metric(
        "test",
        "desc",
        "",
        MetricDescriptorType.GAUGE_DOUBLE,
        label_keys=["a", "b"],
        timeseries=[],
    )
    m.timeseries.append(generate_timeseries(["a", "b"], [1, 2, 3]))
    agent.proxy_export_metrics([m], worker_id_hex=worker_id.hex())
    # Metrics should be exposed.
    assert get_metric(f"{namespace}_test", agent_port) is not None
    agent.clean_all_dead_worker_metrics()
    # Once the worker is dead, metrics should be unavailble.
    assert get_metric(f"{namespace}_test", agent_port) is None
    # Once the worker metrics is re-reported, it is treated as alive again.
    agent.proxy_export_metrics([m], worker_id_hex=worker_id.hex())
    assert get_metric(f"{namespace}_test", agent_port) is not None
    # Clean it again and the worker metrics is cleaned again.
    agent.clean_all_dead_worker_metrics()
    assert get_metric(f"{namespace}_test", agent_port) is None


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_agent_proxy_record_and_export_from_workers_complicated(
    get_agent,
):  # noqa
    """
    Test the complicated worker death case.
    """
    namespace = "test"
    agent, agent_port = get_agent

    # Each worker will report 2 metrics.
    # i.e.,
    # worker 1 => test_1, test_2.
    # worker 2 => test_3, test_4.
    # ...
    worker_ids = [WorkerID.from_random() for _ in range(4)]

    metrics = []
    for i in range(8):
        m = generate_protobuf_metric(
            f"test_{i}",
            "desc",
            "",
            MetricDescriptorType.GAUGE_DOUBLE,
            label_keys=["a", "b"],
            timeseries=[],
        )
        m.timeseries.append(generate_timeseries(["a", str(i)], [3]))
        metrics.append(m)

    i = 0
    for worker_id in worker_ids:
        agent.proxy_export_metrics(
            [metrics[i], metrics[i + 1]], worker_id_hex=worker_id.hex()
        )
        i += 2

    # All metrics must be available.
    for i in range(len(metrics)):
        assert get_metric(f"{namespace}_test_{i}", agent_port) is not None

    # Mark the worker as dead and make sure metrics are properly cleaned.
    i = 0

    while len(worker_ids):
        for worker_id in worker_ids:
            agent.clean_all_dead_worker_metrics()
            assert get_metric(f"{namespace}_test_{i}", agent_port) is None
            assert get_metric(f"{namespace}_test_{i+1}", agent_port) is None

        worker_ids.pop(0)
        metrics.pop(0)
        metrics.pop(0)

        i = 0
        for worker_id in worker_ids:
            agent.proxy_export_metrics(
                [metrics[i], metrics[i + 1]], worker_id_hex=worker_id.hex()
            )
            i += 2

        # Make sure the rest of metrics are still there because new metrics
        # are reported.
        for j in range(i + 2, len(metrics)):
            assert get_metric(f"{namespace}_test_{j}", agent_port) is not None, j
        i += 2


DELAY = 3


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
@pytest.mark.parametrize("get_agent", [DELAY], indirect=True)
def test_metrics_agent_proxy_record_and_export_from_workers_delay(get_agent):  # noqa
    """
    Test the worker metrics are deleted after the delay.
    """
    namespace = "test"
    agent, agent_port = get_agent
    worker_id = WorkerID.from_random()

    m = generate_protobuf_metric(
        "test",
        "desc",
        "",
        MetricDescriptorType.GAUGE_DOUBLE,
        label_keys=["a", "b"],
        timeseries=[],
    )
    m.timeseries.append(generate_timeseries(["a", "b"], [1, 2, 3]))
    agent.proxy_export_metrics([m], worker_id_hex=worker_id.hex())
    agent.clean_all_dead_worker_metrics()
    start = time.time()

    def verify():
        agent.clean_all_dead_worker_metrics()
        return get_metric(f"{namespace}_test", agent_port) is None

    wait_for_condition(verify)
    assert time.time() - start > DELAY


@pytest.mark.skipif(sys.platform == "win32", reason="Flaky on Windows.")
def test_metrics_agent_export_format_correct(get_agent):
    """
    Verifies that there is one metric per metric name and not one
    per metric name + tag combination.
    Also verifies that the prometheus output is in the right format.
    """
    namespace = "test"
    agent, agent_port = get_agent

    # Record a new gauge.
    metric_name = "test"
    test_gauge = Gauge(metric_name, "desc", "unit", ["tag"])
    record_a = Record(
        gauge=test_gauge,
        value=3,
        tags={"tag": "a"},
    )
    agent.record_and_export([record_a])

    # Record a different tag.
    record_b = Record(
        gauge=test_gauge,
        value=4,
        tags={"tag": "b"},
    )
    agent.record_and_export([record_b])

    # Record more than 1 gauge.
    metric_name_2 = "test2"
    test_gauge_2 = Gauge(metric_name_2, "desc", "unit", ["tag"])
    record_c = Record(
        gauge=test_gauge_2,
        value=1,
        tags={"tag": "c"},
    )
    agent.record_and_export([record_c])

    # Basic assertions
    name, samples = get_metric(
        get_prom_metric_name(namespace, metric_name_2), agent_port
    )
    assert name == get_prom_metric_name(namespace, metric_name_2)
    assert len(samples) == 1
    assert samples[0].value == 1
    assert samples[0].labels == {"tag": "c"}

    name, samples = get_metric(get_prom_metric_name(namespace, metric_name), agent_port)
    assert name == get_prom_metric_name(namespace, metric_name)
    assert len(samples) == 2
    assert samples[0].value == 3
    assert samples[0].labels == {"tag": "a"}
    assert samples[1].value == 4
    assert samples[1].labels == {"tag": "b"}

    # Assert there is not multiple HELP text per metric
    # Need to manually parse the prometheus output because the official
    # `prometheus_client.parser` is more lenient than the actual
    # specification and ignores the multiple HELP / TYPE comments.
    metrics_page = "localhost:{}".format(agent_port)
    _, response = list(fetch_raw_prometheus([metrics_page]))[0]
    assert response.count("# HELP test_test desc") == 1
    assert response.count("# TYPE test_test gauge") == 1
    assert response.count("# HELP test_test2 desc") == 1
    assert response.count("# TYPE test_test2 gauge") == 1


def _stub_node_level_metric(label: str, value: float) -> OpencensusProxyMetric:
    metric = OpencensusProxyMetric(
        name="test_metric_01",
        desc="",
        unit="",
        label_keys=["NodeId"],
    )
    metric.add_data(
        (label,),
        LastValueAggregationData(ValueDouble, value),
    )
    return metric


def _stub_worker_level_metric(label: str, value: float) -> OpencensusProxyMetric:
    metric = OpencensusProxyMetric(
        name="test_metric_01",
        desc="",
        unit="",
        label_keys=["NodeId", WORKER_ID_TAG_KEY],
    )
    metric.add_data(
        (label, "worker_01"),
        LastValueAggregationData(ValueDouble, value),
    )
    return metric


def test_aggregate_metric_data():
    collector = OpenCensusProxyCollector("")
    collector._aggregate_metric_data(
        [
            LastValueAggregationData(ValueDouble, 1.0),
            LastValueAggregationData(ValueDouble, 2.0),
            LastValueAggregationData(ValueDouble, 3.0),
        ]
    ).value == 6.0
    collector._aggregate_metric_data(
        [
            SumAggregationData(ValueDouble, 1.0),
            SumAggregationData(ValueDouble, 4.0),
        ]
    ).sum_data == 5.0
    collector._aggregate_metric_data(
        [
            CountAggregationData(1),
            CountAggregationData(1),
        ]
    ).count_data == 2
    with pytest.raises(ValueError, match="Unsupported aggregation type"):
        collector._aggregate_metric_data(
            [
                DistributionAggregationData(
                    mean_data=1.0,
                    count_data=1,
                    sum_of_sqd_deviations=1.0,
                )
            ]
        )


def test_collect_worker_metrics_with_recommended_cardinality():
    aggregated_metrics = OpenCensusProxyCollector(
        ""
    )._aggregate_with_recommended_cardinality(
        [
            _stub_worker_level_metric("node_01", 1.0),
            _stub_worker_level_metric("node_01", 3.0),
            _stub_worker_level_metric("node_02", 2.0),
        ]
    )
    assert len(aggregated_metrics) == 1
    metric = aggregated_metrics[0]
    assert metric.name == "test_metric_01"
    assert metric.label_keys == ["NodeId"]
    # Check that the worker id is removed from the label keys, and the correct metric
    # values are returned.
    assert metric._data.get(("node_01",)).value == 4.0
    assert metric._data.get(("node_02",)).value == 2.0


def test_collect_node_metrics_with_recommended_cardinality():
    aggregated_metrics = OpenCensusProxyCollector(
        ""
    )._aggregate_with_recommended_cardinality(
        [
            _stub_node_level_metric("node_01", 1.0),
            _stub_node_level_metric("node_02", 2.0),
        ]
    )
    # Metrics are already at node level, so they should be returned as is.
    assert len(aggregated_metrics) == 2


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