from os import urandom
from itertools import count
from base64 import b64encode, b64decode
from hashlib import sha256


class pow_base:
    @property
    def base(self):
        return self._base

    @property
    def difficulty(self):
        return self._difficulty

    def __init__(self, difficulty, base=None):
        self._difficulty = difficulty
        self._base = urandom(self._bits//8) if not base else base
        self._target = 2 ** (self._bits - difficulty)

    def work(self, max_attempts=0):
        cycles = count() if max_attempts == 0 else range(max_attempts)
        for nonce in cycles:
            if self.verify(nonce):
                return nonce
        raise RuntimeError('Failed to find a nonce in {} iterations!'.format(max_attempts))

    def verify(self, nonce):
        hexresult = self._hexdigest(self._base + str(nonce).encode())
        return int(hexresult, 16) < self._target

    def export(self):
        return str(self.difficulty) + ':' + b64encode(self.base).decode()

    @classmethod
    def load(cls, export):
        difficulty, base = export.split(':', 1)
        return cls(int(difficulty), b64decode(base.encode()))


class pow_hashlib_base(pow_base):
    def _hexdigest(self, data):
        return self._hasher(data).hexdigest()


class pow256(pow_hashlib_base):
    _bits = 256
    _hasher = sha256