import time
from unittest.mock import MagicMock

import psutil

from metaflow_extensions.observability.plugins.samplers.disk import (
    DiskSampler,
    DiskStats,
)


def _make_io_counters(read_bytes=0, write_bytes=0):
    counters = MagicMock()
    counters.read_bytes = read_bytes
    counters.write_bytes = write_bytes
    return counters


def _make_mock_process(read_bytes=0, write_bytes=0):
    proc = MagicMock(spec=psutil.Process)
    proc.io_counters.return_value = _make_io_counters(read_bytes, write_bytes)
    proc.children.return_value = []
    return proc


class TestDiskSampler:
    def test_returns_zero_stats_when_unavailable(self):
        proc = MagicMock(spec=psutil.Process)
        proc.io_counters.side_effect = AttributeError("not supported on this platform")

        sampler = DiskSampler(process=proc, interval=60.0)
        sampler.start()
        stats = sampler.stop()

        assert isinstance(stats, DiskStats)
        assert stats.read_bytes == 0
        assert stats.write_bytes == 0
        assert stats.avg_read_throughput_mbps == 0.0
        assert stats.avg_write_throughput_mbps == 0.0

    def test_returns_zero_stats_when_access_denied(self):
        proc = MagicMock(spec=psutil.Process)
        proc.io_counters.side_effect = psutil.AccessDenied(pid=1)

        sampler = DiskSampler(process=proc, interval=60.0)
        sampler.start()
        stats = sampler.stop()

        assert stats.read_bytes == 0

    def test_accumulates_read_and_write_bytes(self):
        proc = _make_mock_process()
        # Each call advances the counters: baseline, then two delta snapshots
        proc.io_counters.side_effect = [
            _make_io_counters(read_bytes=0, write_bytes=0),  # baseline
            _make_io_counters(read_bytes=10_000_000, write_bytes=5_000_000),  # +10MB, +5MB
            _make_io_counters(read_bytes=20_000_000, write_bytes=10_000_000),  # +10MB, +5MB
        ]

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

        assert stats.read_bytes >= 10_000_000
        assert stats.write_bytes >= 5_000_000

    def test_aggregates_child_process_io(self):
        parent = _make_mock_process()
        child = _make_mock_process()

        # _probe_availability() consumes one parent read before the baseline.
        parent.io_counters.side_effect = [
            _make_io_counters(0, 0),  # probe
            _make_io_counters(0, 0),  # baseline
            _make_io_counters(5_000_000, 0),  # sample
        ]
        child.io_counters.side_effect = [
            _make_io_counters(0, 0),  # baseline (inside _read_io_counters)
            _make_io_counters(5_000_000, 0),  # sample
        ]
        parent.children.return_value = [child]

        sampler = DiskSampler(process=parent, interval=0.01, include_children=True)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        # Both parent and child contribute
        assert stats.read_bytes >= 10_000_000

    def test_excludes_children_when_disabled(self):
        parent = _make_mock_process()
        child = _make_mock_process()
        parent.io_counters.side_effect = [
            _make_io_counters(0, 0),
            _make_io_counters(4_000_000, 0),
        ]
        child.io_counters.side_effect = [
            _make_io_counters(0, 0),
            _make_io_counters(4_000_000, 0),
        ]
        parent.children.return_value = [child]

        sampler = DiskSampler(process=parent, interval=0.01, include_children=False)
        sampler.start()
        time.sleep(0.05)
        sampler.stop()

        child.io_counters.assert_not_called()

    def test_skips_sample_if_child_disappears(self):
        parent = _make_mock_process()
        child = MagicMock(spec=psutil.Process)
        child.io_counters.side_effect = psutil.NoSuchProcess(pid=999)
        parent.children.return_value = [child]
        parent.io_counters.side_effect = [
            _make_io_counters(0, 0),
            _make_io_counters(1_000_000, 0),
        ]

        sampler = DiskSampler(process=parent, interval=0.01, include_children=True)
        sampler.start()
        time.sleep(0.05)
        stats = sampler.stop()

        # Should not crash; parent IO still counted
        assert isinstance(stats, DiskStats)
