import hashlib
import sys

def check_difficulty(hash_bytes: bytes, difficulty_bits: int) -> bool:
    bits = difficulty_bits
    i = 0
    while bits > 0:
        byte = hash_bytes[i]
        if bits >= 8:
            if byte != 0:
                return False
            bits -= 8
        else:
            mask = 0xFF << (8 - bits) & 0xFF
            if (byte & mask) != 0:
                return False
            bits = 0
        i += 1
    return True

def main():
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <challenge_string> <difficulty_bits>")
        sys.exit(1)

    challenge = sys.argv[1]
    try:
        difficulty_bits = int(sys.argv[2])
    except ValueError:
        print("Difficulty must be a number")
        sys.exit(1)

    if difficulty_bits <= 0 or difficulty_bits > 256:
        print("Difficulty bits must be between 1 and 256")
        sys.exit(1)

    counter = 0
    while True:
        input_str = f"{challenge}{counter}"
        hash_bytes = hashlib.sha256(input_str.encode()).digest()
        if check_difficulty(hash_bytes, difficulty_bits):
            print("====================")
            print(f"Found nonce: {counter}")
            print("====================")
            print(f"Hash: {hash_bytes.hex()}")
            break
        if counter & 0xFFFFF == 0:
            print(f"Attempts: {counter}")
        counter += 1

if __name__ == "__main__":
    main()