from __future__ import annotations

import ctypes
import json
import math
import os
import threading
import time
import uuid
from collections.abc import Mapping
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Any, Self

from models import MetricRecord, RunContext, SystemTelemetryRecord
from storage import LocalSessionStore, utc_now

_MIRRORED_METRIC_SPECS = (
    ("step_time_ms", "system/runtime/step_time_ms", "Wall-clock duration for a training step."),
    ("data_time_ms", "system/runtime/data_time_ms", "Estimated data-loading duration for a step."),
    ("compute_time_ms", "system/runtime/compute_time_ms", "Estimated compute duration for a step."),
    ("cpu_usage_percent", "system/cpu/usage_percent", "CPU usage sampled during runtime telemetry."),
    ("ram_usage_percent", "system/ram/usage_percent", "RAM usage percentage sampled during runtime telemetry."),
    ("ram_used_bytes", "system/ram/used_bytes", "RAM used in bytes sampled during runtime telemetry."),
    ("process_rss_bytes", "system/process/rss_bytes", "Process resident memory sampled during runtime telemetry."),
    (
        "gpu_utilization_percent",
        "system/gpu/utilization_percent",
        "GPU utilization sampled during runtime telemetry.",
    ),
    (
        "gpu_memory_used_bytes",
        "system/gpu/memory_used_bytes",
        "GPU memory used in bytes sampled during runtime telemetry.",
    ),
    (
        "gpu_memory_usage_percent",
        "system/gpu/memory_usage_percent",
        "GPU memory usage percentage sampled during runtime telemetry.",
    ),
    (
        "overflow_count",
        "system/runtime/overflow_count",
        "Cumulative overflow count observed by the telemetry module.",
    ),
    (
        "skipped_step_count",
        "system/runtime/skipped_step_count",
        "Cumulative skipped-step count observed by the telemetry module.",
    ),
)


@dataclass(frozen=True, slots=True)
class SystemTelemetryConfig:
    """Configuration for step sampling cadence and optional machine probes."""

    sample_every_n_steps: int = 1
    emit_scalar_metrics: bool = True
    capture_cpu: bool = True
    capture_gpu: bool = True
    gpu_index: int | None = None


@dataclass(frozen=True, slots=True)
class _CpuSnapshot:
    backend: str
    usage_percent: float | None = None
    fallback_reason: str | None = None
    ram_used_bytes: int | None = None
    ram_total_bytes: int | None = None
    ram_available_bytes: int | None = None
    ram_usage_percent: float | None = None
    process_rss_bytes: int | None = None


@dataclass(frozen=True, slots=True)
class _GpuSnapshot:
    backend: str
    index: int | None = None
    name: str | None = None
    utilization_percent: float | None = None
    memory_used_bytes: int | None = None
    memory_total_bytes: int | None = None
    memory_usage_percent: float | None = None
    fallback_reason: str | None = None


def _import_psutil() -> Any | None:
    try:
        import psutil  # type: ignore[import-not-found]
    except Exception:
        return None
    return psutil


def _import_pynvml() -> Any | None:
    try:
        import pynvml  # type: ignore[import-not-found]
    except Exception:
        return None
    return pynvml


def _import_torch() -> Any | None:
    try:
        import torch  # type: ignore[import-not-found]
    except Exception:
        return None
    return torch


class _UnavailableCpuSampler:
    def __init__(self, fallback_reason: str) -> None:
        self._fallback_reason = fallback_reason

    def sample(self) -> _CpuSnapshot:
        return _CpuSnapshot(backend="none", fallback_reason=self._fallback_reason)

    def close(self) -> None:
        return None


class _PsutilCpuSampler:
    def __init__(self, psutil_module: Any) -> None:
        self._psutil = psutil_module
        self._process = psutil_module.Process(os.getpid())
        self._psutil.cpu_percent(interval=None)

    def sample(self) -> _CpuSnapshot:
        memory = self._psutil.virtual_memory()
        process_memory = self._process.memory_info()
        return _CpuSnapshot(
            backend="psutil",
            usage_percent=float(self._psutil.cpu_percent(interval=None)),
            ram_used_bytes=int(memory.used),
            ram_total_bytes=int(memory.total),
            ram_available_bytes=int(memory.available),
            ram_usage_percent=float(memory.percent),
            process_rss_bytes=int(process_memory.rss),
        )

    def close(self) -> None:
        return None


class _StdlibCpuSampler:
    def __init__(self, fallback_reason: str) -> None:
        self._fallback_reason = fallback_reason
        self._last_wall_time = time.perf_counter()
        self._last_process_time = time.process_time()
        self._cpu_count = max(os.cpu_count() or 1, 1)

    def sample(self) -> _CpuSnapshot:
        wall_time = time.perf_counter()
        process_time = time.process_time()
        elapsed_wall = max(wall_time - self._last_wall_time, 0.0)
        elapsed_process = max(process_time - self._last_process_time, 0.0)

        self._last_wall_time = wall_time
        self._last_process_time = process_time

        usage_percent: float | None = 0.0
        if elapsed_wall > 0.0:
            usage_percent = min(100.0, (elapsed_process / elapsed_wall) * 100.0 / self._cpu_count)

        total_bytes, available_bytes = _get_system_memory_bytes()
        used_bytes = None
        usage_ratio = None
        if total_bytes is not None and available_bytes is not None and total_bytes > 0:
            used_bytes = max(total_bytes - available_bytes, 0)
            usage_ratio = (used_bytes / total_bytes) * 100.0

        return _CpuSnapshot(
            backend="process",
            usage_percent=usage_percent,
            fallback_reason=self._fallback_reason,
            ram_used_bytes=used_bytes,
            ram_total_bytes=total_bytes,
            ram_available_bytes=available_bytes,
            ram_usage_percent=usage_ratio,
            process_rss_bytes=_get_process_rss_bytes(),
        )

    def close(self) -> None:
        return None


class _UnavailableGpuSampler:
    def __init__(self, fallback_reason: str) -> None:
        self._fallback_reason = fallback_reason

    def sample(self) -> _GpuSnapshot:
        return _GpuSnapshot(backend="none", fallback_reason=self._fallback_reason)

    def close(self) -> None:
        return None


class _NvmlGpuSampler:
    def __init__(self, pynvml_module: Any, gpu_index: int | None) -> None:
        self._pynvml = pynvml_module
        self._pynvml.nvmlInit()
        device_count = int(self._pynvml.nvmlDeviceGetCount())
        if device_count <= 0:
            self._pynvml.nvmlShutdown()
            raise RuntimeError("NVML did not report any available GPU devices.")

        resolved_index = 0 if gpu_index is None else int(gpu_index)
        if resolved_index < 0 or resolved_index >= device_count:
            self._pynvml.nvmlShutdown()
            raise ValueError(
                f"Requested GPU index {resolved_index} is outside the available NVML range 0..{device_count - 1}."
            )

        self._gpu_index = resolved_index
        self._handle = self._pynvml.nvmlDeviceGetHandleByIndex(resolved_index)

    def sample(self) -> _GpuSnapshot:
        utilization = self._pynvml.nvmlDeviceGetUtilizationRates(self._handle)
        memory = self._pynvml.nvmlDeviceGetMemoryInfo(self._handle)
        gpu_name = self._pynvml.nvmlDeviceGetName(self._handle)
        if isinstance(gpu_name, bytes):
            gpu_name = gpu_name.decode("utf-8", errors="replace")
        usage_percent = None
        if int(memory.total) > 0:
            usage_percent = (int(memory.used) / int(memory.total)) * 100.0
        return _GpuSnapshot(
            backend="pynvml",
            index=self._gpu_index,
            name=str(gpu_name),
            utilization_percent=float(utilization.gpu),
            memory_used_bytes=int(memory.used),
            memory_total_bytes=int(memory.total),
            memory_usage_percent=usage_percent,
        )

    def close(self) -> None:
        try:
            self._pynvml.nvmlShutdown()
        except Exception:
            pass


class _TorchGpuSampler:
    def __init__(self, torch_module: Any, gpu_index: int | None, fallback_reason: str) -> None:
        self._torch = torch_module
        self._fallback_reason = fallback_reason
        cuda = getattr(torch_module, "cuda", None)
        if cuda is None or not bool(cuda.is_available()):
            raise RuntimeError("torch.cuda is not available.")

        device_count = int(cuda.device_count())
        if device_count <= 0:
            raise RuntimeError("torch.cuda did not report any GPU devices.")

        resolved_index = 0 if gpu_index is None else int(gpu_index)
        if resolved_index < 0 or resolved_index >= device_count:
            raise ValueError(
                f"Requested GPU index {resolved_index} is outside the available CUDA range 0..{device_count - 1}."
            )
        self._gpu_index = resolved_index

    def sample(self) -> _GpuSnapshot:
        cuda = self._torch.cuda
        total_bytes = None
        used_bytes = None
        if hasattr(cuda, "mem_get_info"):
            free_bytes, total_bytes = cuda.mem_get_info(self._gpu_index)
            used_bytes = int(total_bytes) - int(free_bytes)
            total_bytes = int(total_bytes)
        else:
            used_bytes = int(cuda.memory_allocated(self._gpu_index))
            if hasattr(cuda, "get_device_properties"):
                properties = cuda.get_device_properties(self._gpu_index)
                total_bytes = int(getattr(properties, "total_memory", 0) or 0) or None

        usage_percent = None
        if total_bytes is not None and total_bytes > 0 and used_bytes is not None:
            usage_percent = (used_bytes / total_bytes) * 100.0

        return _GpuSnapshot(
            backend="torch_cuda",
            index=self._gpu_index,
            name=str(cuda.get_device_name(self._gpu_index)),
            utilization_percent=None,
            memory_used_bytes=used_bytes,
            memory_total_bytes=total_bytes,
            memory_usage_percent=usage_percent,
            fallback_reason=self._fallback_reason,
        )

    def close(self) -> None:
        return None


def _build_cpu_sampler(*, capture_cpu: bool) -> Any:
    if not capture_cpu:
        return _UnavailableCpuSampler("CPU telemetry disabled by config.")

    psutil_module = _import_psutil()
    if psutil_module is not None:
        try:
            return _PsutilCpuSampler(psutil_module)
        except Exception as exc:
            return _StdlibCpuSampler(f"psutil unavailable at runtime ({exc}); using process fallback.")

    return _StdlibCpuSampler("psutil unavailable; using process fallback.")


def _build_gpu_sampler(*, capture_gpu: bool, gpu_index: int | None) -> Any:
    if not capture_gpu:
        return _UnavailableGpuSampler("GPU telemetry disabled by config.")

    nvml_failure_reason: str | None = None
    pynvml_module = _import_pynvml()
    if pynvml_module is not None:
        try:
            return _NvmlGpuSampler(pynvml_module, gpu_index)
        except Exception as exc:
            nvml_failure_reason = f"pynvml unavailable at runtime ({exc})."
    else:
        nvml_failure_reason = "pynvml unavailable."

    torch_module = _import_torch()
    if torch_module is not None:
        try:
            return _TorchGpuSampler(
                torch_module,
                gpu_index=gpu_index,
                fallback_reason=f"{nvml_failure_reason} Falling back to torch.cuda memory telemetry.",
            )
        except Exception as exc:
            return _UnavailableGpuSampler(f"{nvml_failure_reason} torch.cuda unavailable ({exc}).")

    return _UnavailableGpuSampler(f"{nvml_failure_reason} torch.cuda unavailable.")


def _get_system_memory_bytes() -> tuple[int | None, int | None]:
    if os.name == "nt":
        return _get_windows_system_memory_bytes()

    try:
        page_size = int(os.sysconf("SC_PAGE_SIZE"))
        total_pages = int(os.sysconf("SC_PHYS_PAGES"))
        available_pages = int(os.sysconf("SC_AVPHYS_PAGES"))
    except (AttributeError, OSError, ValueError):
        return None, None

    return page_size * total_pages, page_size * available_pages


def _get_windows_system_memory_bytes() -> tuple[int | None, int | None]:
    class _MemoryStatusEx(ctypes.Structure):
        _fields_ = [
            ("dwLength", ctypes.c_ulong),
            ("dwMemoryLoad", ctypes.c_ulong),
            ("ullTotalPhys", ctypes.c_ulonglong),
            ("ullAvailPhys", ctypes.c_ulonglong),
            ("ullTotalPageFile", ctypes.c_ulonglong),
            ("ullAvailPageFile", ctypes.c_ulonglong),
            ("ullTotalVirtual", ctypes.c_ulonglong),
            ("ullAvailVirtual", ctypes.c_ulonglong),
            ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
        ]

    memory_status = _MemoryStatusEx()
    memory_status.dwLength = ctypes.sizeof(_MemoryStatusEx)
    try:
        success = ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(memory_status))
    except Exception:
        return None, None
    if not success:
        return None, None
    return int(memory_status.ullTotalPhys), int(memory_status.ullAvailPhys)


def _get_process_rss_bytes() -> int | None:
    if os.name == "nt":
        return _get_windows_process_rss_bytes()

    try:
        import resource
    except Exception:
        return None

    try:
        usage = resource.getrusage(resource.RUSAGE_SELF)
    except Exception:
        return None

    rss = int(usage.ru_maxrss)
    if rss <= 0:
        return None

    if hasattr(os, "uname") and os.uname().sysname == "Darwin":
        return rss
    return rss * 1024


def _get_windows_process_rss_bytes() -> int | None:
    class _ProcessMemoryCounters(ctypes.Structure):
        _fields_ = [
            ("cb", ctypes.c_ulong),
            ("PageFaultCount", ctypes.c_ulong),
            ("PeakWorkingSetSize", ctypes.c_size_t),
            ("WorkingSetSize", ctypes.c_size_t),
            ("QuotaPeakPagedPoolUsage", ctypes.c_size_t),
            ("QuotaPagedPoolUsage", ctypes.c_size_t),
            ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t),
            ("QuotaNonPagedPoolUsage", ctypes.c_size_t),
            ("PagefileUsage", ctypes.c_size_t),
            ("PeakPagefileUsage", ctypes.c_size_t),
        ]

    counters = _ProcessMemoryCounters()
    counters.cb = ctypes.sizeof(_ProcessMemoryCounters)
    try:
        process_handle = ctypes.windll.kernel32.GetCurrentProcess()
        success = ctypes.windll.psapi.GetProcessMemoryInfo(
            process_handle,
            ctypes.byref(counters),
            counters.cb,
        )
    except Exception:
        return None
    if not success:
        return None
    return int(counters.WorkingSetSize)


@dataclass(slots=True)
class StepTelemetryScope:
    """Mutable helper that measures one training step with optional data/compute boundaries."""

    _telemetry: "SystemTelemetry"
    step: int
    epoch: int | None
    _started_at: float
    _base_metadata: dict[str, Any] | None = None
    _data_ready_at: float | None = None
    _finished: bool = False

    def mark_data_ready(self) -> None:
        """Mark the end of data loading and the start of compute for the tracked step."""

        if self._finished or self._data_ready_at is not None:
            return
        self._data_ready_at = self._telemetry._clock()

    def finish(
        self,
        *,
        overflow_occurred: bool = False,
        skipped_step: bool = False,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> SystemTelemetryRecord | None:
        """Finalize the step, capture machine state, and queue a telemetry record."""

        if self._finished:
            raise RuntimeError("This step telemetry scope has already been finished.")

        finished_at = self._telemetry._clock()
        step_time_ms = max(finished_at - self._started_at, 0.0) * 1000.0
        data_time_ms = None
        compute_time_ms = None
        if self._data_ready_at is not None:
            data_time_ms = max(self._data_ready_at - self._started_at, 0.0) * 1000.0
            compute_time_ms = max(finished_at - self._data_ready_at, 0.0) * 1000.0

        merged_metadata = dict(self._base_metadata or {})
        if metadata is not None:
            merged_metadata.update(dict(metadata))

        self._finished = True
        return self._telemetry.record_step(
            self.step,
            epoch=self.epoch,
            step_time_ms=step_time_ms,
            data_time_ms=data_time_ms,
            compute_time_ms=compute_time_ms,
            overflow_occurred=overflow_occurred,
            skipped_step=skipped_step,
            metadata=merged_metadata or None,
            timestamp=timestamp,
        )

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        if not self._finished:
            self.finish()


class SystemTelemetry:
    """Buffered runtime telemetry writer aligned to training steps or wall-clock boundaries."""

    def __init__(
        self,
        session: Any,
        *,
        config: SystemTelemetryConfig | None = None,
        batch_size: int = 20,
        _cpu_sampler: Any | None = None,
        _gpu_sampler: Any | None = None,
        _clock: Any | None = None,
        _timestamp_factory: Any | None = None,
    ) -> None:
        if batch_size <= 0:
            raise ValueError("batch_size must be a positive integer.")

        self._context, self._store = self._resolve_session_dependencies(session)
        self._session = session if callable(getattr(session, "_register_flushable", None)) else None
        self._config = config or SystemTelemetryConfig()
        if self._config.sample_every_n_steps <= 0:
            raise ValueError("SystemTelemetryConfig.sample_every_n_steps must be positive.")

        self._batch_size = batch_size
        self._buffer: list[SystemTelemetryRecord] = []
        self._metric_buffer: list[MetricRecord] = []
        self._lock = threading.RLock()
        self._closed = False
        self._clock = _clock or time.perf_counter
        self._timestamp_factory = _timestamp_factory or utc_now
        self._cpu_sampler = _cpu_sampler or _build_cpu_sampler(capture_cpu=self._config.capture_cpu)
        self._gpu_sampler = _gpu_sampler or _build_gpu_sampler(
            capture_gpu=self._config.capture_gpu,
            gpu_index=self._config.gpu_index,
        )
        self._overflow_count, self._skipped_step_count = self._load_existing_counters()
        if self._session is not None:
            self._session._register_flushable(self)

    @property
    def run_id(self) -> str:
        return self._context.run_id

    @property
    def buffered_count(self) -> int:
        with self._lock:
            return len(self._buffer)

    def start_step(
        self,
        step: int,
        *,
        epoch: int | None = None,
        metadata: Mapping[str, Any] | None = None,
    ) -> StepTelemetryScope:
        """Start a step-aligned telemetry scope that can later record data and compute timing."""

        self._ensure_open()
        return StepTelemetryScope(
            _telemetry=self,
            step=self._normalize_index("step", step, required=True),
            epoch=self._normalize_index("epoch", epoch),
            _started_at=self._clock(),
            _base_metadata=self._normalize_metadata(metadata),
        )

    def sample(
        self,
        *,
        step: int | None = None,
        epoch: int | None = None,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> SystemTelemetryRecord:
        """Capture a point-in-time telemetry sample aligned to either a step or wall-clock time."""

        self._ensure_open()
        normalized_step = self._normalize_index("step", step)
        normalized_epoch = self._normalize_index("epoch", epoch)
        normalized_metadata = self._normalize_metadata(metadata)
        telemetry_record = self._build_record(
            boundary="step" if normalized_step is not None else "time",
            step=normalized_step,
            epoch=normalized_epoch,
            step_time_ms=None,
            data_time_ms=None,
            compute_time_ms=None,
            overflow_occurred=False,
            skipped_step=False,
            overflow_count=self._overflow_count,
            skipped_step_count=self._skipped_step_count,
            metadata=normalized_metadata,
            timestamp=timestamp,
        )
        self._enqueue(telemetry_record)
        return telemetry_record

    def record_step(
        self,
        step: int,
        *,
        epoch: int | None = None,
        step_time_ms: float | None = None,
        data_time_ms: float | None = None,
        compute_time_ms: float | None = None,
        overflow_occurred: bool = False,
        skipped_step: bool = False,
        metadata: Mapping[str, Any] | None = None,
        timestamp: str | None = None,
    ) -> SystemTelemetryRecord | None:
        """Queue step-aligned telemetry from explicit timing inputs."""

        self._ensure_open()
        normalized_step = self._normalize_index("step", step, required=True)
        normalized_epoch = self._normalize_index("epoch", epoch)
        normalized_step_time = self._normalize_duration("step_time_ms", step_time_ms)
        normalized_data_time = self._normalize_duration("data_time_ms", data_time_ms)
        normalized_compute_time = self._normalize_duration("compute_time_ms", compute_time_ms)
        normalized_metadata = self._normalize_metadata(metadata)
        overflow_flag = bool(overflow_occurred)
        skipped_flag = bool(skipped_step)

        if normalized_step_time is None:
            if normalized_data_time is None and normalized_compute_time is None:
                raise ValueError(
                    "record_step requires step_time_ms or at least one of data_time_ms / compute_time_ms."
                )
            normalized_step_time = float((normalized_data_time or 0.0) + (normalized_compute_time or 0.0))

        with self._lock:
            if overflow_flag:
                self._overflow_count += 1
            if skipped_flag:
                self._skipped_step_count += 1
            overflow_count = self._overflow_count
            skipped_step_count = self._skipped_step_count

        if not self._should_sample_step(normalized_step, overflow_flag, skipped_flag):
            return None

        telemetry_record = self._build_record(
            boundary="step",
            step=normalized_step,
            epoch=normalized_epoch,
            step_time_ms=normalized_step_time,
            data_time_ms=normalized_data_time,
            compute_time_ms=normalized_compute_time,
            overflow_occurred=overflow_flag,
            skipped_step=skipped_flag,
            overflow_count=overflow_count,
            skipped_step_count=skipped_step_count,
            metadata=normalized_metadata,
            timestamp=timestamp,
        )
        self._enqueue(telemetry_record)
        return telemetry_record

    def flush(self) -> tuple[SystemTelemetryRecord, ...]:
        """Persist buffered telemetry records and mirrored scalar metrics in one transaction."""

        self._ensure_open()
        with self._lock:
            if not self._buffer and not self._metric_buffer:
                return ()
            telemetry_batch = tuple(self._buffer)
            mirrored_metrics = tuple(self._metric_buffer)
            self._buffer.clear()
            self._metric_buffer.clear()

        try:
            return self._store.write_system_telemetry(
                self.run_id,
                telemetry_batch,
                scalar_metrics=mirrored_metrics,
            )
        except Exception:
            with self._lock:
                self._buffer = list(telemetry_batch) + self._buffer
                self._metric_buffer = list(mirrored_metrics) + self._metric_buffer
            raise

    def read(
        self,
        *,
        boundary: str | None = None,
        step: int | None = None,
        step_min: int | None = None,
        step_max: int | None = None,
        epoch: int | None = None,
        epoch_min: int | None = None,
        epoch_max: int | None = None,
        limit: int | None = None,
    ) -> tuple[SystemTelemetryRecord, ...]:
        """Flush pending telemetry and return persisted records ordered by sequence."""

        if not self._closed:
            self.flush()
        return self._store.read_system_telemetry(
            self.run_id,
            boundary=boundary,
            step=step,
            step_min=step_min,
            step_max=step_max,
            epoch=epoch,
            epoch_min=epoch_min,
            epoch_max=epoch_max,
            limit=limit,
        )

    def close(self) -> None:
        """Flush buffered telemetry and release any underlying machine probes."""

        if self._closed:
            return
        try:
            self.flush()
        finally:
            self._closed = True
            if self._session is not None:
                self._session._unregister_flushable(self)
            self._safe_close_sampler(self._cpu_sampler)
            self._safe_close_sampler(self._gpu_sampler)

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        self.close()

    def _enqueue(self, record: SystemTelemetryRecord) -> None:
        mirrored_metrics = self._build_mirrored_metrics(record)
        should_flush = False
        with self._lock:
            self._buffer.append(record)
            self._metric_buffer.extend(mirrored_metrics)
            should_flush = len(self._buffer) >= self._batch_size

        if should_flush:
            self.flush()

    def _build_record(
        self,
        *,
        boundary: str,
        step: int | None,
        epoch: int | None,
        step_time_ms: float | None,
        data_time_ms: float | None,
        compute_time_ms: float | None,
        overflow_occurred: bool,
        skipped_step: bool,
        overflow_count: int,
        skipped_step_count: int,
        metadata: dict[str, Any] | None,
        timestamp: str | None,
    ) -> SystemTelemetryRecord:
        cpu_snapshot = self._sample_cpu()
        gpu_snapshot = self._sample_gpu()
        return SystemTelemetryRecord(
            id=uuid.uuid4().hex,
            run_id=self.run_id,
            sequence_id=self._store.allocate_system_telemetry_sequence(self.run_id),
            boundary=boundary,
            step=step,
            epoch=epoch,
            timestamp=timestamp or self._timestamp_factory(),
            step_time_ms=step_time_ms,
            data_time_ms=data_time_ms,
            compute_time_ms=compute_time_ms,
            cpu_backend=cpu_snapshot.backend,
            cpu_usage_percent=cpu_snapshot.usage_percent,
            cpu_fallback_reason=cpu_snapshot.fallback_reason,
            ram_used_bytes=cpu_snapshot.ram_used_bytes,
            ram_total_bytes=cpu_snapshot.ram_total_bytes,
            ram_available_bytes=cpu_snapshot.ram_available_bytes,
            ram_usage_percent=cpu_snapshot.ram_usage_percent,
            process_rss_bytes=cpu_snapshot.process_rss_bytes,
            gpu_backend=gpu_snapshot.backend,
            gpu_index=gpu_snapshot.index,
            gpu_name=gpu_snapshot.name,
            gpu_utilization_percent=gpu_snapshot.utilization_percent,
            gpu_memory_used_bytes=gpu_snapshot.memory_used_bytes,
            gpu_memory_total_bytes=gpu_snapshot.memory_total_bytes,
            gpu_memory_usage_percent=gpu_snapshot.memory_usage_percent,
            gpu_fallback_reason=gpu_snapshot.fallback_reason,
            overflow_occurred=overflow_occurred,
            step_skipped=skipped_step,
            overflow_count=overflow_count,
            skipped_step_count=skipped_step_count,
            metadata=metadata,
        )

    def _build_mirrored_metrics(self, record: SystemTelemetryRecord) -> tuple[MetricRecord, ...]:
        if not self._config.emit_scalar_metrics:
            return ()
        if record.step is None and record.epoch is None:
            return ()

        metrics: list[MetricRecord] = []
        for attribute_name, metric_key, description in _MIRRORED_METRIC_SPECS:
            value = getattr(record, attribute_name)
            if value is None or isinstance(value, bool):
                continue
            metrics.append(
                MetricRecord(
                    id=uuid.uuid4().hex,
                    run_id=self.run_id,
                    sequence_id=self._store.allocate_metric_sequence(self.run_id),
                    namespace=self._extract_namespace(metric_key),
                    metric_name=self._extract_metric_name(metric_key),
                    metric_key=metric_key,
                    value=float(value),
                    step=record.step,
                    epoch=record.epoch,
                    timestamp=record.timestamp,
                    description=description,
                    metadata={
                        "source": "system_telemetry",
                        "boundary": record.boundary,
                        "cpu_backend": record.cpu_backend,
                        "gpu_backend": record.gpu_backend,
                    },
                )
            )
        return tuple(metrics)

    def _sample_cpu(self) -> _CpuSnapshot:
        try:
            snapshot = self._cpu_sampler.sample()
        except Exception as exc:
            return _CpuSnapshot(
                backend="none",
                fallback_reason=f"CPU telemetry sampling failed: {exc}",
            )
        return snapshot if isinstance(snapshot, _CpuSnapshot) else _CpuSnapshot(backend="none")

    def _sample_gpu(self) -> _GpuSnapshot:
        try:
            snapshot = self._gpu_sampler.sample()
        except Exception as exc:
            return _GpuSnapshot(
                backend="none",
                fallback_reason=f"GPU telemetry sampling failed: {exc}",
            )
        return snapshot if isinstance(snapshot, _GpuSnapshot) else _GpuSnapshot(backend="none")

    def _load_existing_counters(self) -> tuple[int, int]:
        try:
            existing_records = self._store.read_system_telemetry(self.run_id)
        except Exception:
            return 0, 0
        if not existing_records:
            return 0, 0
        latest = existing_records[-1]
        return latest.overflow_count, latest.skipped_step_count

    def _should_sample_step(self, step: int, overflow_occurred: bool, skipped_step: bool) -> bool:
        if overflow_occurred or skipped_step:
            return True
        return step % self._config.sample_every_n_steps == 0

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("SystemTelemetry is closed.")

    @staticmethod
    def _safe_close_sampler(sampler: Any) -> None:
        try:
            sampler.close()
        except Exception:
            pass

    @staticmethod
    def _resolve_session_dependencies(session: Any) -> tuple[RunContext, LocalSessionStore]:
        if isinstance(session, RunContext):
            return session, LocalSessionStore(session.storage_dir)

        context = getattr(session, "context", None)
        store = getattr(session, "_store", None)
        if isinstance(context, RunContext) and isinstance(store, LocalSessionStore):
            return context, store

        raise TypeError(
            "SystemTelemetry expects a TrainingSession or RunContext created by the Quark SDK."
        )

    @staticmethod
    def _normalize_index(label: str, value: int | None, *, required: bool = False) -> int | None:
        if value is None:
            if required:
                raise ValueError(f"{label} is required.")
            return None
        if isinstance(value, bool) or not isinstance(value, Integral):
            raise TypeError(f"{label} must be an integer when provided.")
        normalized_value = int(value)
        if normalized_value < 0:
            raise ValueError(f"{label} must not be negative.")
        return normalized_value

    @staticmethod
    def _normalize_duration(label: str, value: float | None) -> float | None:
        if value is None:
            return None
        if isinstance(value, bool) or not isinstance(value, Real):
            raise TypeError(f"{label} must be a numeric scalar when provided.")
        normalized_value = float(value)
        if not math.isfinite(normalized_value):
            raise ValueError(f"{label} must be finite when provided.")
        if normalized_value < 0.0:
            raise ValueError(f"{label} must not be negative.")
        return normalized_value

    @staticmethod
    def _normalize_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any] | None:
        if metadata is None:
            return None
        normalized = dict(metadata)
        json.dumps(normalized, sort_keys=True)
        return normalized

    @staticmethod
    def _extract_namespace(metric_key: str) -> str:
        if "/" not in metric_key:
            return ""
        return metric_key.rsplit("/", 1)[0]

    @staticmethod
    def _extract_metric_name(metric_key: str) -> str:
        if "/" not in metric_key:
            return metric_key
        return metric_key.rsplit("/", 1)[1]
