"""
==========================
vipin Performance Counters
==========================

Structs for counting performance metrics.

"""
from __future__ import annotations

import json
from time import time
from typing import TYPE_CHECKING

import psutil

if TYPE_CHECKING:
    from psutil import scpufreq, scpustats, sdiskio, snetio  # type: ignore[attr-defined]


class CounterSnapshot:
    def __init__(
        self,
        cpu: scpustats | None = None,
        disk: sdiskio | None = None,
        freq: scpufreq | None = None,
        net: snetio | None = None,
        timestamp: float | None = None,
    ) -> None:
        self.cpu = psutil.cpu_stats() if cpu is None else cpu
        self.disk = (
            psutil.disk_io_counters(perdisk=False, nowrap=True) if disk is None else disk
        )
        self.freq = psutil.cpu_freq(percpu=False) if freq is None else freq
        self.net = psutil.net_io_counters(pernic=False, nowrap=True) if net is None else net
        self.timestamp = time() if timestamp is None else timestamp

    def to_dict(self) -> dict[str, dict[str, int | float] | float]:
        c_dict: dict[str, dict[str, int | float] | float] = {}
        # Handle potential None values from psutil calls
        if self.cpu is not None:
            c_dict["cpu"] = dict(self.cpu._asdict())
        if self.freq is not None:
            c_dict["freq"] = dict(self.freq._asdict())
        if self.disk is not None:
            c_dict["disk"] = dict(self.disk._asdict())
        if self.net is not None:
            c_dict["net"] = dict(self.net._asdict())
        c_dict["time"] = self.timestamp
        return c_dict

    def __sub__(self, other: "CounterSnapshot") -> "CounterSnapshot":
        # Handle potential None values and use Any for psutil types
        cpu = None
        if self.cpu is not None and other.cpu is not None:
            cpu = type(other.cpu)(
                *tuple(self.cpu[i] - other.cpu[i] for i in range(len(other.cpu)))
            )

        disk = None
        if self.disk is not None and other.disk is not None:
            disk = type(other.disk)(
                *tuple(self.disk[i] - other.disk[i] for i in range(len(other.disk)))
            )

        freq = None
        if self.freq is not None and other.freq is not None:
            freq = type(other.freq)(
                *tuple((self.freq[i] + other.freq[i]) / 2 for i in range(len(other.freq)))
            )

        net = None
        if self.net is not None and other.net is not None:
            net = type(other.net)(
                *tuple(self.net[i] - other.net[i] for i in range(len(other.net)))
            )

        timestamp = self.timestamp - other.timestamp
        return CounterSnapshot(cpu, disk, freq, net, timestamp)

    def __repr__(self) -> str:
        return json.dumps(self.to_dict())
