from __future__ import annotations

import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"

if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from quark import SystemTelemetry, SystemTelemetryConfig, TrainingSession
from system_telemetry import (
    _CpuSnapshot,
    _GpuSnapshot,
    _UnavailableGpuSampler,
)
from tests.service_test_utils import unique_project_name


class _StaticCpuSampler:
    def __init__(self, snapshot: _CpuSnapshot) -> None:
        self._snapshot = snapshot

    def sample(self) -> _CpuSnapshot:
        return self._snapshot

    def close(self) -> None:
        return None


class _StaticGpuSampler:
    def __init__(self, snapshot: _GpuSnapshot) -> None:
        self._snapshot = snapshot

    def sample(self) -> _GpuSnapshot:
        return self._snapshot

    def close(self) -> None:
        return None


class _SequenceClock:
    def __init__(self, values: list[float]) -> None:
        self._values = list(values)

    def __call__(self) -> float:
        if not self._values:
            raise RuntimeError("Clock exhausted.")
        return self._values.pop(0)


class _SequenceTimestamps:
    def __init__(self, values: list[str]) -> None:
        self._values = list(values)

    def __call__(self) -> str:
        if not self._values:
            raise RuntimeError("Timestamp factory exhausted.")
        return self._values.pop(0)


class _FakeCuda:
    @staticmethod
    def is_available() -> bool:
        return True

    @staticmethod
    def device_count() -> int:
        return 1

    @staticmethod
    def mem_get_info(index: int) -> tuple[int, int]:
        return 400, 1000

    @staticmethod
    def get_device_name(index: int) -> str:
        return "Fake RTX"


class _FakeTorch:
    cuda = _FakeCuda()


class SystemTelemetryTests(unittest.TestCase):
    def test_step_scope_persists_step_aligned_telemetry_and_mirrored_timing_metrics(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("telemetry-demo"),
                storage_dir=storage_dir,
            )
            telemetry = SystemTelemetry(
                session,
                batch_size=10,
                _cpu_sampler=_StaticCpuSampler(
                    _CpuSnapshot(
                        backend="psutil",
                        usage_percent=41.5,
                        ram_used_bytes=1024,
                        ram_total_bytes=2048,
                        ram_available_bytes=1024,
                        ram_usage_percent=50.0,
                        process_rss_bytes=512,
                    )
                ),
                _gpu_sampler=_StaticGpuSampler(
                    _GpuSnapshot(
                        backend="pynvml",
                        index=0,
                        name="Unit Test GPU",
                        utilization_percent=72.0,
                        memory_used_bytes=800,
                        memory_total_bytes=1000,
                        memory_usage_percent=80.0,
                    )
                ),
                _clock=_SequenceClock([0.0, 0.015, 0.055]),
                _timestamp_factory=_SequenceTimestamps(["2026-04-14T10:00:00Z"]),
            )

            scope = telemetry.start_step(4, epoch=1, metadata={"phase": "train"})
            scope.mark_data_ready()
            record = scope.finish()

            self.assertIsNotNone(record)
            assert record is not None
            self.assertEqual(record.boundary, "step")
            self.assertEqual(record.step, 4)
            self.assertEqual(record.epoch, 1)
            self.assertAlmostEqual(record.step_time_ms or 0.0, 55.0)
            self.assertAlmostEqual(record.data_time_ms or 0.0, 15.0)
            self.assertAlmostEqual(record.compute_time_ms or 0.0, 40.0)
            self.assertEqual(record.cpu_backend, "psutil")
            self.assertEqual(record.gpu_backend, "pynvml")
            self.assertEqual(record.metadata, {"phase": "train"})

            persisted_records = telemetry.read()
            self.assertEqual(len(persisted_records), 1)
            self.assertAlmostEqual(persisted_records[0].step_time_ms or 0.0, 55.0)

            step_time_metrics = session._store.read_scalar_metrics(
                session.run_id,
                metric_key="system/runtime/step_time_ms",
            )
            data_time_metrics = session._store.read_scalar_metrics(
                session.run_id,
                metric_key="system/runtime/data_time_ms",
            )
            self.assertEqual(len(step_time_metrics), 1)
            self.assertEqual(len(data_time_metrics), 1)
            self.assertAlmostEqual(step_time_metrics[0].value, 55.0)
            self.assertAlmostEqual(data_time_metrics[0].value, 15.0)

    def test_cpu_fallback_uses_process_sampler_when_psutil_is_unavailable(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("cpu-fallback"),
                storage_dir=storage_dir,
            )

            with patch("system_telemetry._import_psutil", return_value=None):
                telemetry = SystemTelemetry(
                    session,
                    _gpu_sampler=_UnavailableGpuSampler("GPU telemetry unavailable for this test."),
                    _timestamp_factory=_SequenceTimestamps(["2026-04-14T10:00:01Z"]),
                )
                telemetry.sample(step=0)
                record = telemetry.read()[0]

            self.assertEqual(record.cpu_backend, "process")
            self.assertIsNotNone(record.cpu_fallback_reason)
            self.assertIn("psutil unavailable", record.cpu_fallback_reason or "")
            self.assertEqual(record.gpu_backend, "none")

    def test_gpu_fallback_uses_torch_cuda_when_nvml_is_missing(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("gpu-fallback"),
                storage_dir=storage_dir,
            )

            with patch("system_telemetry._import_pynvml", return_value=None), patch(
                "system_telemetry._import_torch",
                return_value=_FakeTorch(),
            ):
                telemetry = SystemTelemetry(
                    session,
                    _cpu_sampler=_StaticCpuSampler(_CpuSnapshot(backend="none")),
                    _timestamp_factory=_SequenceTimestamps(["2026-04-14T10:00:02Z"]),
                )
                telemetry.record_step(0, step_time_ms=12.5)
                record = telemetry.read()[0]

            self.assertEqual(record.gpu_backend, "torch_cuda")
            self.assertEqual(record.gpu_name, "Fake RTX")
            self.assertEqual(record.gpu_memory_used_bytes, 600)
            self.assertEqual(record.gpu_memory_total_bytes, 1000)
            self.assertAlmostEqual(record.gpu_memory_usage_percent or 0.0, 60.0)
            self.assertIsNone(record.gpu_utilization_percent)
            self.assertIn("pynvml unavailable", record.gpu_fallback_reason or "")


if __name__ == "__main__":
    unittest.main(verbosity=2)
