"""System metrics collector for heartbeat reporting.

Collects CPU, memory, and execution metrics for the heartbeat message
as specified in Section 5.3 (heartbeat schema) and Section 14.
"""

from __future__ import annotations

import logging
import time
from typing import Any

import psutil

logger = logging.getLogger("qremotex.metrics")


class MetricsCollector:
    """Collects system and execution metrics."""

    def __init__(self) -> None:
        self._start_time = time.time()
        self._total_executions: int = 0
        self._total_errors: int = 0
        self._active_executions: int = 0
        self._queued_executions: int = 0

    # ------------------------------------------------------------------ #
    # Execution counters
    # ------------------------------------------------------------------ #
    def record_execution_start(self) -> None:
        self._active_executions += 1

    def record_execution_end(self, success: bool) -> None:
        self._active_executions = max(0, self._active_executions - 1)
        self._total_executions += 1
        if not success:
            self._total_errors += 1

    def set_queued_count(self, count: int) -> None:
        self._queued_executions = count

    # ------------------------------------------------------------------ #
    # Snapshot for heartbeat
    # ------------------------------------------------------------------ #
    def snapshot(self) -> dict[str, Any]:
        """Return current metrics snapshot for the heartbeat message."""
        mem = psutil.virtual_memory()
        return {
            "cpu_percent": psutil.cpu_percent(interval=None),
            "memory_used_mb": round(mem.used / (1024**2)),
            "memory_total_mb": round(mem.total / (1024**2)),
            "active_executions": self._active_executions,
            "queued_executions": self._queued_executions,
            "uptime_seconds": int(time.time() - self._start_time),
            "total_executions": self._total_executions,
            "total_errors": self._total_errors,
        }
