# -*- coding: utf-8 -*-
#
#       SystemInformation.py
#
#       "THE BEER-WARE LICENSE" (Revision 42):
#       Dionisio E Alonso <dealonso@gmail.com> wrote this file. As long as you
#       retain this notice you can do whatever you want with this stuff. If we
#       meet some day, and you think this stuff is worth it, you can buy me a
#       beer in return
#
#        with  by Baco
#

import winreg
import ctypes
import string


def get_registry_value(key, subkey, value):
    key = getattr(winreg, key)
    handle = winreg.OpenKey(key, subkey)
    (value, type) = winreg.QueryValueEx(handle, value)
    return value


class SystemInformation:
    def __init__(self):
        self.cpu = self._cpu().strip()
        self.graphics = self._graphics()
        self.total_ram, self.available_ram = self._ram()
        self.hdd_free = self._disk()
        self.os = self._os_version()

    def _cpu(self):
        return get_registry_value(
            'HKEY_LOCAL_MACHINE',
            'HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0',
            'ProcessorNameString')

    def _graphics(self):
        reg_keys = set()
        for device in range(16):
            try:
                video_card_string = get_registry_value(
                    'HKEY_LOCAL_MACHINE',
                    'HARDWARE\\DEVICEMAP\\VIDEO',
                    '\Device\Video{}'.format(device))
                reg_key = video_card_string.split('\\')[-2]
                # Probe keys to avoid software devices
                get_registry_value(
                    'HKEY_LOCAL_MACHINE',
                    'SYSTEM\\CurrentControlSet\\Control\\Video\\'
                    '{}\\0000'.format(reg_key),
                    'HardwareInformation.AdapterString')
                reg_keys.add(reg_key)
            except FileNotFoundError:
                pass

        video_cards = []
        for video_card in reg_keys:
            try:
                device_descrption = get_registry_value(
                    'HKEY_LOCAL_MACHINE',
                    'SYSTEM\\CurrentControlSet\\Control\\Video\\'
                    '{}\\0000'.format(video_card),
                    'Device Description')
            except FileNotFoundError:
                try:
                    device_descrption = get_registry_value(
                        'HKEY_LOCAL_MACHINE',
                        'SYSTEM\\CurrentControlSet\\Control\\Video\\'
                        '{}\\0000'.format(video_card),
                        'DriverDesc')
                except FileNotFoundError:
                    device_descrption = 'Failed to obtain video adapter.'
            video_cards.append(device_descrption)
        return '\n\t'.join(video_cards)

    def _ram(self):
        kernel32 = ctypes.windll.kernel32
        c_ulong = ctypes.c_ulong

        class MemoryStatus(ctypes.Structure):
            _fields_ = [('dwLength', c_ulong),
                        ('dwMemoryLoad', c_ulong),
                        ('dwTotalPhys', c_ulong),
                        ('dwAvailPhys', c_ulong),
                        ('dwTotalPageFile', c_ulong),
                        ('dwAvailPageFile', c_ulong),
                        ('dwTotalVirtual', c_ulong),
                        ('dwAvailVirtual', c_ulong)]

        mem_status = MemoryStatus()
        mem_status.dwLength = ctypes.sizeof(MemoryStatus)
        kernel32.GlobalMemoryStatus(ctypes.byref(mem_status))
        return (mem_status.dwTotalPhys, mem_status.dwAvailPhys)

    def _disk(self):
        drives = [c + ':' for c in list(string.ascii_uppercase[2:])]
        freeuser = ctypes.c_int64()
        total = ctypes.c_int64()
        free = ctypes.c_int64()
        free_values = []
        for drive in drives:
            ctypes.windll.kernel32.GetDiskFreeSpaceExW(drive,
                                                       ctypes.byref(freeuser),
                                                       ctypes.byref(total),
                                                       ctypes.byref(free))
            if freeuser.value != 0:
                free_values.append((drive, freeuser.value))
            freeuser.value = 0
        return free_values

    def _os_version(self):
        def get(key):
            return get_registry_value(
                'HKEY_LOCAL_MACHINE',
                'SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion', key)
        os = get('ProductName')
        try:
            sp = get('CSDVersion')
        except FileNotFoundError:
            try:
                sp = 'n' + get('CSDBuildNumber')
            except FileNotFoundError:
                sp = '(no service pack found)'
        build = get('CurrentBuildNumber')
        return '{} {} (build {})'.format(os, sp, build)

    def save(self):
        # hdd_gb = [(drive, free/(1024**3)) for drive, free in self.hdd_free]
        # part_strs = ['{} {:.2f}GB free'.format(*part) for part in hdd_gb]
        part_strs = []
        for p in self.hdd_free:
            part_strs.append('{} {:.2f}GB free'.format(p[0], p[1]/(1024**3)))
        hdd_string = ', '.join(part_strs)
        with open('system_information.txt', 'w') as f:
            f.write('System Information\n'
                    '==================\n\n')
            f.write('CPU:\t{}\n'.format(self.cpu))
            f.write('Video:\t{}\n'.format(self.graphics))
            f.write('RAM:\t{:.2f}GB total\n'.format(self.total_ram/(1024**3)))
            f.write('HDD:\t{}\n'.format(hdd_string))
            f.write('OS:\t{}\n'.format(self.os))

    def __str__(self):
        part_strs = []
        for p in self.hdd_free:
            part_strs.append('{} {:.2f}GB free'.format(p[0], p[1]/(1024**3)))
        hdd_string = ', '.join(part_strs)
        return '\n'.join(['CPU:\t{}'.format(self.cpu),
                          'Video:\t{}'.format(self.graphics),
                          'RAM:\t{:.2f}GB total'
                          .format(self.total_ram/(1024**3)),
                          'HDD:\t{}'.format(hdd_string),
                          'OS:\t{}'.format(self.os)])

if __name__ == '__main__':
    sys = SystemInformation()
    sys.save()
    print(sys)