import csv
import re

from libqtile.widget import base

sensors_mapping = {
    'fan_speed': 'fan.speed',
    'perf': 'pstate',
    'temp': 'temperature.gpu',
}


def _all_sensors_names_correct(sensors):
    return all(map(lambda x: x in sensors_mapping, sensors))


class NvidiaSensors(base.ThreadPoolText):
    """Displays temperature, fan speed and performance level Nvidia GPU."""

    orientations = base.ORIENTATION_HORIZONTAL
    defaults = [
        (
            'format', '{temp}°C',
            'Display string format. Three options available:  '
            '``{temp}`` - temperature, ``{fan_speed}`` and ``{perf}`` - '
            'performance level'
        ),
        ('foreground_alert', 'ff0000', 'Foreground colour alert'),
        (
            'gpu_bus_id', '',
            "GPU's Bus ID, ex: ``01:00.0``. If leave empty will display all "
            "available GPU's"
        ),
        ('update_interval', 2, 'Update interval in seconds.'),
        (
            'threshold',
            70,
            'If the current temperature value is above, '
            'then change to foreground_alert colour'
        ),
    ]

    def __init__(self, **config):
        base.ThreadPoolText.__init__(self, '', **config)
        self.add_defaults(NvidiaSensors.defaults)
        self.foreground_normal = self.foreground

    def _get_sensors_data(self, command):
        return csv.reader(
            self.call_process(
                command, shell=True
            ).strip().replace(' ', '').split('\n')
        )

    def _parse_format_string(self):
        return {sensor for sensor in re.findall('{(.+?)}', self.format)}

    def poll(self):
        sensors = self._parse_format_string()
        if not _all_sensors_names_correct(sensors):
            return 'Wrong sensor name'
        bus_id = f'-i {self.gpu_bus_id}' if self.gpu_bus_id else ''
        command = 'nvidia-smi {} --query-gpu={} --format=csv,noheader'.format(
            bus_id, ",".join(sensors_mapping[sensor] for sensor in sensors)
        )
        try:
            sensors_data = [
                dict(zip(sensors, gpu))
                for gpu in self._get_sensors_data(command)
            ]
            for gpu in sensors_data:
                if gpu.get('temp'):
                    if int(gpu['temp']) > self.threshold:
                        self.foreground = self.foreground_alert
                    else:
                        self.foreground = self.foreground_normal
            return ' - '.join(
                [self.format.format(**gpu) for gpu in sensors_data]
            )
        except Exception:
            return None
