import winreg
from collections import namedtuple

__all__ = ['Reg', 'REG_ROOTS', 'query_reg']

_empty_generator = (_ele for _ele in ())
_backslash = '\\'

_REG_VALUE = namedtuple("REG_VALUE", ('name', 'value', 'data_type'))

REG_PERMISSION_DESC = [
    'KEY_ALL_ACCESS',
    'KEY_WRITE',
    'KEY_SET_VALUE',
    'KEY_EXECUTE',
    'KEY_WOW64_32KEY',
    'KEY_WOW64_64KEY',
    'KEY_CREATE_LINK',
    'KEY_CREATE_SUB_KEY',
    'KEY_ENUMERATE_SUB_KEYS',
    'KEY_READ',
    'KEY_QUERY_VALUE',
    'KEY_NOTIFY'
]

REG_PERMISSION = namedtuple("REG_PERMISSION", REG_PERMISSION_DESC)(*(
    winreg.KEY_ALL_ACCESS,
    winreg.KEY_WRITE,
    winreg.KEY_SET_VALUE,
    winreg.KEY_EXECUTE,
    winreg.KEY_WOW64_32KEY,
    winreg.KEY_WOW64_64KEY,
    winreg.KEY_CREATE_LINK,
    winreg.KEY_CREATE_SUB_KEY,
    winreg.KEY_ENUMERATE_SUB_KEYS,
    winreg.KEY_READ,
    winreg.KEY_QUERY_VALUE,
    winreg.KEY_NOTIFY
))

REG_ROOTS_PATH = (
    'HKEY_CLASSES_ROOT',
    'HKEY_CURRENT_USER',
    'HKEY_LOCAL_MACHINE',
    'HKEY_USERS',
    'HKEY_CURRENT_CONFIG'
)

REG_ROOTS = namedtuple('REG_ROOTS', REG_ROOTS_PATH)(*[
    winreg.HKEY_CLASSES_ROOT,
    winreg.HKEY_CURRENT_USER,
    winreg.HKEY_LOCAL_MACHINE,
    winreg.HKEY_USERS,
    winreg.HKEY_CURRENT_CONFIG
])


class Reg:
    def __init__(self, reg_key, subkey='', parent=None):
        self._parent = parent
        if not self._parent:
            if isinstance(reg_key, int) and REG_ROOTS.count(reg_key) > 0:
                self._parent = REG_ROOTS_PATH[REG_ROOTS.index(reg_key)]
        self._key = f"{self._parent}{_backslash if subkey else ''}{subkey.encode('utf-8', 'ignore').decode('utf-8', 'ignore')}"
        self.keyname = subkey
        self._self = None
        self._permission = None
        self._permission_desc = None
        for permission in REG_PERMISSION:
            try:
                self._self = winreg.OpenKey(reg_key, subkey, 0, permission)
                self._permission = permission
                self._permission_desc = REG_PERMISSION_DESC[REG_PERMISSION.index(permission)]
                break
            except Exception:
                continue

    def __str__(self):
        return f"Reg<\"{self._key}\", {self._permission_desc}>"

    __repr__ = __str__

    def __iter__(self):
        if self.has_keys:
            for _key in self.keys:
                yield _key
        if self.has_values:
            for _value in self.values:
                yield _value

    @property
    def keys(self):
        if not self._self:
            return _empty_generator
        _count = 0
        while True:
            try:
                yield Reg(self._self, winreg.EnumKey(self._self, _count), self._key)
            except WindowsError:
                break
            _count += 1

    @property
    def values(self):
        if not self._self:
            return _empty_generator
        _count = 0
        while True:
            try:
                yield _REG_VALUE(*winreg.EnumValue(self._self, _count))
            except WindowsError:
                break
            _count += 1

    @property
    def has_values(self):
        try:
            self.values.__next__()
        except StopIteration:
            return False
        return True

    @property
    def has_keys(self):
        try:
            self.keys.__next__()
        except StopIteration:
            return False
        return True

    def delete(self, warning=True):
        """删除注册表键值，包含子键值"""
        if self.has_keys:
            for subkey in self.keys:
                subkey.delete()
        winreg.DeleteKeyEx(self._self, sub_key='', access=self._permission)


def query_reg(reg: Reg, *query_texts: str):
    assert isinstance(reg, Reg)
    if query_texts is None or len(query_texts) == 0:
        return _empty_generator
    for _mainitem in reg:
        if isinstance(_mainitem, Reg):
            if query_texts.count(_mainitem.keyname) > 0:
                yield _mainitem
            if _mainitem.has_keys:
                for _subitem in query_reg(_mainitem, *query_texts):
                    yield _subitem
        if isinstance(_mainitem, _REG_VALUE):
            if query_texts.count(_mainitem.name) > 0:
                yield _mainitem


if __name__ == '__main__':
    def test_query():
        for _item in query_reg(Reg(REG_ROOTS.HKEY_LOCAL_MACHINE), 'Google'):
            print(_item)
            if isinstance(_item, Reg):
                print(list(_item.values))


    for item in Reg(REG_ROOTS.HKEY_CLASSES_ROOT):
        print(item)

    # test_query()