import collections, os.path
import _winreg

# Alternative to support unicode in Python 2
# http://pypi.python.org/pypi/winreg_unicode
#import winreg_unicode as _winreg 

for key in dir(_winreg):
    if key.startswith('HKEY_'):
        globals()[key[5:]] = getattr(_winreg, key)
del key

def interpret(value_type, value):
    if value_type == _winreg.REG_SZ:
        return value
    elif value_type == _winreg.REG_EXPAND_SZ:
        return os.path.expandvars(value)
    else:
        return None

class HiveKeyIterator(collections.KeysView):
    def __init__(self, hivekey):
        self._key = _winreg.OpenKey(hivekey._hive, hivekey._key)
        self._index = 0
        self._done = False

    def next(self):
        if self._done:
            raise StopIteration
        while True:
            for i in range(3):
                try:
                    name, data, data_type = _winreg.EnumValue(self._key,
                                                              self._index)
                    break
                except WindowsError:
                    continue
            else:
                self._done = True
                self._key.Close()
                raise StopIteration

            self._index += 1
            rv = interpret(data_type, name)
            if rv is not None:
                return rv

class HiveKey(collections.MutableMapping):
    def __init__(self, hive, key=''):
        self._hive = hive
        self._key = key

    def _getvalue(self, key, subkey):
        with _winreg.OpenKey(self._hive, key) as root_key:
            value, value_type = _winreg.QueryValueEx(root_key, subkey)
        rv = interpret(value_type, value)
        if rv is None:
            raise NotImplementedError('type: %d' % value_type)
        return rv

    def __getitem__(self, key):
        key, subkey = self._compute_subkey(key)
        try:
            return self._getvalue(key, subkey)
        except WindowsError:
            pass
        key = key + '\\' + subkey
        try:
            with _winreg.OpenKey(self._hive, key): # verify that key exists
                pass
        except WindowsError:
            raise KeyError
        return HiveKey(self._hive, key)

    def __iter__(self):
        return HiveKeyIterator(self)

    def __len__(self):
        "value is approximate since .keys() drops non-string types"
        with _winreg.OpenKey(self._hive, self._key) as k:
            num_subkeys, num_values, t = _winreg.QueryInfoKey(k)
            return num_values

    def _compute_subkey(self, k):
        key = (self._key + '\\' + k.strip('\\')).strip('\\')
        subkeys = key.split('\\')
        key = '\\'.join(subkeys[:-1])
        subkey = subkeys[-1]
        return key, subkey

    def __setitem__(self, k, v):
        if type(v) != type(u''):
            raise NotImplementedError
        key, subkey = self._compute_subkey(k)
        with _winreg.CreateKey(self._hive, key) as root_key:
            _winreg.SetValueEx(root_key, subkey, 0, _winreg.REG_SZ, v)

    def __delitem__(self, k):
        key, subkey = self._compute_subkey(k)
        with _winreg.OpenKey(self._hive, key,
                             0, _winreg.KEY_ALL_ACCESS) as root_key:
            try:
                _winreg.DeleteValue(root_key, subkey)
            except WindowsError:
                try:
                    _winreg.DeleteKey(root_key, subkey)
                except WindowsError:
                    raise KeyError

    def get_keys(self):
        values = []
        with _winreg.OpenKey(self._hive, self._key) as root_key:
            try:
                for i in range(100000):
                    key = _winreg.EnumKey(root_key, i)
                    if '\x00' in key: continue
                    values.append(key)
            except WindowsError:
                pass
        return values