from __future__ import annotations

import os
import socket
import time
from datetime import datetime, timezone
from pathlib import Path
from types import ModuleType

import psutil

from device_gateway.protocol import ServerSnapshot


class SystemCollector:
    """Keep only the latest local snapshot and one network counter frame."""

    def __init__(
        self,
        disk_path: str | Path = "/",
        *,
        hostname: str | None = None,
        psutil_module: ModuleType = psutil,
    ) -> None:
        self.disk_path = str(disk_path)
        self.hostname = hostname or socket.gethostname()
        self.psutil = psutil_module
        self.latest: ServerSnapshot | None = None
        self._previous_network: tuple[float, int, int] | None = None

    def collect(self) -> ServerSnapshot:
        now_monotonic = time.monotonic()
        memory = self.psutil.virtual_memory()
        swap = self.psutil.swap_memory()
        disk = self.psutil.disk_usage(self.disk_path)
        network = self.psutil.net_io_counters()

        rx_bps = 0.0
        tx_bps = 0.0
        if self._previous_network is not None:
            previous_time, previous_rx, previous_tx = self._previous_network
            elapsed = now_monotonic - previous_time
            if elapsed > 0:
                rx_bps = max(0.0, (network.bytes_recv - previous_rx) / elapsed)
                tx_bps = max(0.0, (network.bytes_sent - previous_tx) / elapsed)
        self._previous_network = (now_monotonic, network.bytes_recv, network.bytes_sent)

        load_1m, load_5m, _ = os.getloadavg()
        snapshot = ServerSnapshot(
            hostname=self.hostname,
            cpu_percent=self.psutil.cpu_percent(interval=None),
            load_1m=load_1m,
            load_5m=load_5m,
            memory_used=memory.used,
            memory_total=memory.total,
            memory_percent=memory.percent,
            swap_percent=swap.percent,
            disk_used=disk.used,
            disk_total=disk.total,
            disk_percent=disk.percent,
            network_rx_bytes=network.bytes_recv,
            network_tx_bytes=network.bytes_sent,
            network_rx_bps=rx_bps,
            network_tx_bps=tx_bps,
            uptime=max(0, int(time.time() - self.psutil.boot_time())),
            timestamp=datetime.now(timezone.utc).isoformat(),
        )
        self.latest = snapshot
        return snapshot