import time
from unittest.mock import MagicMock

import psutil
import pytest

from metaflow_extensions.observability.plugins.samplers.cpu import (
    CpuMemorySampler,
    CpuMemoryStats,
)

_RSS_BYTES = 200 * 1024 * 1024  # 200 MB


def _make_mock_process(cpu_percent=50.0, rss=_RSS_BYTES):
    proc = MagicMock(spec=psutil.Process)
    proc.cpu_percent.return_value = cpu_percent
    mem_info = MagicMock()
    mem_info.rss = rss
    proc.memory_info.return_value = mem_info
    return proc


class TestCpuMemorySampler:
    def test_returns_zero_stats_when_no_samples_collected(self):
        proc = _make_mock_process()
        sampler = CpuMemorySampler(process=proc, interval=60.0)
        sampler.start()
        stats = sampler.stop()

        assert isinstance(stats, CpuMemoryStats)
        assert stats.avg_cpu_percent == 0.0
        assert stats.max_cpu_percent == 0.0
        assert stats.p95_cpu_percent == 0.0
        assert stats.avg_memory_mb == 0.0
        assert stats.max_memory_mb == 0.0

    def test_primes_cpu_measurement_on_start(self):
        proc = _make_mock_process()
        sampler = CpuMemorySampler(process=proc, interval=60.0)
        sampler.start()
        sampler.stop()

        # cpu_percent() must be called once before sampling starts
        # so the kernel has a reference point for the delta
        assert proc.cpu_percent.call_count >= 1

    def test_computes_avg_from_multiple_samples(self):
        proc = _make_mock_process()
        proc.cpu_percent.side_effect = [0.0, 20.0, 40.0, 60.0]  # prime + 3 samples

        sampler = CpuMemorySampler(process=proc, interval=0.01)
        sampler.start()
        time.sleep(0.08)
        stats = sampler.stop()

        assert stats.avg_cpu_percent == pytest.approx(40.0, abs=5.0)

    def test_computes_max_from_samples(self):
        proc = _make_mock_process()
        proc.cpu_percent.side_effect = [0.0, 10.0, 90.0, 50.0, 50.0, 50.0, 50.0]

        sampler = CpuMemorySampler(process=proc, interval=0.01)
        sampler.start()
        time.sleep(0.08)
        stats = sampler.stop()

        assert stats.max_cpu_percent == pytest.approx(90.0, abs=5.0)

    def test_computes_memory_mb_from_rss(self):
        proc = _make_mock_process(rss=512_000_000)  # exactly 512 metric MB

        sampler = CpuMemorySampler(process=proc, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        assert stats.avg_memory_mb == pytest.approx(512.0, abs=1.0)
        assert stats.max_memory_mb == pytest.approx(512.0, abs=1.0)

    def test_skips_sample_when_process_disappears(self):
        proc = _make_mock_process()
        proc.cpu_percent.side_effect = [
            0.0,  # prime
            psutil.NoSuchProcess(pid=999),
        ]

        sampler = CpuMemorySampler(process=proc, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        # No crash — should return zero stats
        assert isinstance(stats, CpuMemoryStats)

    def test_p95_with_single_sample(self):
        proc = _make_mock_process(cpu_percent=75.0)
        proc.cpu_percent.side_effect = [0.0, 75.0, 75.0, 75.0]

        sampler = CpuMemorySampler(process=proc, interval=0.01)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        assert stats.p95_cpu_percent == pytest.approx(75.0, abs=1.0)
