from dataclasses import dataclass

from psutil import cpu_percent, sensors_temperatures, virtual_memory


@dataclass
class TemperatureSensor:
    sensor_type: str
    sensor_label: str
    current_temp: float
    highest_temp: float | None = None
    critical_temp: float | None = None


@dataclass
class CpuSensor:
    current_load: float
    highest_load: float = 90
    critical_load: float = 95


@dataclass
class RamSensor:
    current_avail: int
    warning_avail: int = 4 * 1024**3
    critical_avail: int = 2 * 1024**3


class Sensors:
    @staticmethod
    def get_temperatures() -> dict[str, list[TemperatureSensor]]:
        psutil_temp_sensors = sensors_temperatures()

        temp_sensors = {}

        for s_type, sensors in psutil_temp_sensors.items():
            if s_type not in temp_sensors.keys():
                temp_sensors[s_type] = []
            match (s_type):
                case "nvme":
                    for sensor in sensors:
                        temp_sensors[s_type].append(
                            TemperatureSensor(
                                sensor_type=s_type,
                                sensor_label=sensor.label,
                                current_temp=sensor.current,
                                highest_temp=sensor.high,
                                critical_temp=sensor.critical,
                            )
                        )
                case "amdgpu":
                    temp_sensors[s_type].append(
                        TemperatureSensor(
                            sensor_type=s_type,
                            sensor_label="Integrated GPU",
                            current_temp=sensors[0].current,
                        )
                    )
                case "k10temp":
                    temp_sensors[s_type].append(
                        TemperatureSensor(
                            sensor_type=s_type,
                            sensor_label="AMD CPU",
                            current_temp=sensors[0].current,
                            critical_temp=95.0,  # hardcoded because we have R9 7900X
                        )
                    )
                case "nct6687":
                    for sensor in sensors:
                        temp_sensors[s_type].append(
                            TemperatureSensor(
                                sensor_type=s_type,
                                sensor_label=sensor.label,
                                current_temp=sensor.current,
                                highest_temp=sensor.high or None,
                                critical_temp=sensor.critical or None,
                            )
                        )

        return temp_sensors

    @staticmethod
    def get_cpu() -> CpuSensor:
        return CpuSensor(current_load=cpu_percent())

    @staticmethod
    def get_ram() -> RamSensor:
        return RamSensor(current_avail=virtual_memory().available)


if __name__ == "__main__":
    for i in Sensors.get_temperatures():
        print(i)