import io
import logging
import re
import time
from unittest.mock import MagicMock, patch

import pytest

import ray
from ray.data._internal.execution.interfaces.physical_operator import (
    OpTask,
    PhysicalOperator,
    RefBundle,
    TaskExecDriverStats,
)
from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry
from ray.data._internal.execution.operators.input_data_buffer import (
    InputDataBuffer,
)
from ray.data._internal.execution.operators.task_pool_map_operator import (
    MapOperator,
)
from ray.data._internal.issue_detection.detectors.hanging_detector import (
    DEFAULT_OP_TASK_STATS_MIN_COUNT,
    DEFAULT_OP_TASK_STATS_STD_FACTOR,
    HangingExecutionIssueDetector,
    HangingExecutionIssueDetectorConfig,
)
from ray.data._internal.issue_detection.detectors.high_memory_detector import (
    HighMemoryIssueDetector,
)
from ray.data._internal.util import GiB
from ray.data.block import BlockMetadata, TaskExecWorkerStats
from ray.data.context import DataContext
from ray.tests.conftest import *  # noqa


class FakeOpTask(OpTask):
    """A fake OpTask for testing purposes."""

    def __init__(self, task_index: int):
        super().__init__(task_index)

    def get_waitable(self):
        """Return a dummy waitable."""
        return ray.put(None)


class FakeOperator(PhysicalOperator):
    def __init__(self, name: str, data_context: DataContext):
        super().__init__(name=name, input_dependencies=[], data_context=data_context)

    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
        pass

    def has_next(self) -> bool:
        return False

    def _get_next_inner(self) -> RefBundle:
        assert False

    def get_stats(self):
        return {}

    def get_active_tasks(self):
        # Return active tasks based on what's in _running_tasks
        # This ensures has_execution_finished() works correctly
        return [FakeOpTask(task_idx) for task_idx in self.metrics._running_tasks]


class TestHangingExecutionIssueDetector:
    def test_hanging_detector_configuration(self, restore_data_context):
        """Test hanging detector configuration and initialization."""
        # Test default configuration from DataContext
        ctx = DataContext.get_current()
        default_config = ctx.issue_detectors_config.hanging_detector_config
        assert default_config.op_task_stats_min_count == DEFAULT_OP_TASK_STATS_MIN_COUNT
        assert (
            default_config.op_task_stats_std_factor == DEFAULT_OP_TASK_STATS_STD_FACTOR
        )

        # Test custom configuration
        min_count = 5
        std_factor = 3.0
        custom_config = HangingExecutionIssueDetectorConfig(
            op_task_stats_min_count=min_count,
            op_task_stats_std_factor=std_factor,
        )
        ctx.issue_detectors_config.hanging_detector_config = custom_config

        detector = HangingExecutionIssueDetector(
            dataset_id="id", operators=[], config=custom_config
        )
        assert detector._op_task_stats_min_count == min_count
        assert detector._op_task_stats_std_factor_threshold == std_factor

    @patch(
        "ray.data._internal.execution.interfaces.op_runtime_metrics.DistributionTracker"
    )
    def test_basic_hanging_detection(
        self, mock_stats_cls, ray_start_regular_shared, restore_data_context
    ):
        # Set up logging capture
        log_capture = io.StringIO()
        handler = logging.StreamHandler(log_capture)
        logger = logging.getLogger("ray.data._internal.issue_detection")
        logger.addHandler(handler)

        # Set up mock stats to return values that will trigger adaptive threshold
        mocked_mean = 2.0  # Increase from 0.5 to 2.0 seconds
        mocked_stddev = 0.2  # Increase from 0.05 to 0.2 seconds
        mock_stats = mock_stats_cls.return_value
        mock_stats.num_samples = 20  # Enough samples
        mock_stats.mean = mocked_mean
        mock_stats.stddev = mocked_stddev

        # Explicitly enable hanging detection for this test
        ctx = DataContext.get_current()
        ctx.issue_detectors_config.detectors = [HangingExecutionIssueDetector]
        detector_cfg = ctx.issue_detectors_config.hanging_detector_config
        detector_cfg.detection_time_interval_s = 0.00

        # test no hanging doesn't log hanging warning
        def f1(x):
            return x

        _ = ray.data.range(1).map(f1).materialize()

        log_output = log_capture.getvalue()
        warn_msg = r"A task \(task_id=.+\) of operator .+(?:\(pid=.+, node_id=.+, attempt=.+\) )?has been running or stuck in scheduling for [\d\.]+s"
        assert re.search(warn_msg, log_output) is None, log_output

        # # test hanging does log hanging warning
        def f2(x):
            time.sleep(5.0)  # Increase from 1.1 to 5.0 seconds to exceed new threshold
            return x

        _ = ray.data.range(1).map(f2).materialize()

        log_output = log_capture.getvalue()
        assert re.search(warn_msg, log_output) is not None, log_output

    @patch("time.perf_counter")
    def test_hanging_detector_detects_issues(
        self, mock_perf_counter, ray_start_regular_shared
    ):
        """Test that the hanging detector correctly identifies tasks that exceed the adaptive threshold."""

        # Configure hanging detector with extreme std_factor values
        config = HangingExecutionIssueDetectorConfig(
            op_task_stats_min_count=1,
            op_task_stats_std_factor=1,
            detection_time_interval_s=0,
        )
        op = FakeOperator("TestOperator", DataContext.get_current())
        detector = HangingExecutionIssueDetector(
            dataset_id="test_dataset", operators=[op], config=config
        )

        # Create a simple RefBundle for testing
        block_ref = ray.put([{"id": 0}])
        metadata = BlockMetadata(
            num_rows=1, size_bytes=1, exec_stats=None, input_files=None
        )
        input_bundle = RefBundle(
            blocks=(BlockEntry(block_ref, metadata),),
            owns_blocks=True,
            schema=None,
        )

        mock_perf_counter.return_value = 0.0

        # Submit three tasks. Two of them finish immediately, while the third one hangs.
        op.metrics.on_task_submitted(0, input_bundle)
        op.metrics.on_task_submitted(1, input_bundle)
        op.metrics.on_task_submitted(2, input_bundle)
        op.metrics.on_task_finished(
            0,
            exception=None,
            task_exec_stats=TaskExecWorkerStats(task_wall_time_s=1.0),
            task_exec_driver_stats=TaskExecDriverStats(task_output_backpressure_s=0),
        )
        op.metrics.on_task_finished(
            1,
            exception=None,
            task_exec_stats=TaskExecWorkerStats(task_wall_time_s=1.0),
            task_exec_driver_stats=TaskExecDriverStats(task_output_backpressure_s=0),
        )

        # Start detecting — all tasks were submitted at t=0, so no time has elapsed.
        issues = detector.detect()
        assert len(issues) == 0

        # Advance perf_counter to trigger the issue detection
        mock_perf_counter.return_value = 10.0

        # On the second detect() call, the hanging task should be detected
        issues = detector.detect()
        assert len(issues) > 0, "Expected hanging issue to be detected"
        assert issues[0].issue_type.value == "hanging"
        assert "has been running or stuck in scheduling for" in issues[0].message
        assert "longer than the average task duration" in issues[0].message


@pytest.mark.parametrize(
    "configured_memory, actual_memory, should_return_issue",
    [
        # User has appropriately configured memory, so no issue.
        (8 * GiB, 8 * GiB, False),
        # User hasn't configured memory correctly and memory use is high, so issue.
        (None, 8 * GiB, True),
        (1 * GiB, 8 * GiB, True),
        # User hasn't configured memory correctly but memory use is low, so no issue.
        (None, 1 * GiB, False),
    ],
)
def test_high_memory_detection(
    configured_memory, actual_memory, should_return_issue, restore_data_context
):
    ctx = DataContext.get_current()

    input_data_buffer = InputDataBuffer(ctx, input_data=[])
    map_operator = MapOperator.create(
        map_transformer=MagicMock(),
        input_op=input_data_buffer,
        data_context=ctx,
        ray_remote_args={"memory": configured_memory},
    )
    map_operator._metrics = MagicMock(average_max_uss_per_task=actual_memory)
    topology = {input_data_buffer: MagicMock(), map_operator: MagicMock()}

    operators = list(topology.keys())

    detector = HighMemoryIssueDetector(
        dataset_id="id",
        operators=operators,
        config=ctx.issue_detectors_config.high_memory_detector_config,
    )
    issues = detector.detect()

    assert should_return_issue == bool(issues)


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))