#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <openssl/sha.h>

int check_difficulty(const unsigned char *hash, int difficulty_bits) {
    int bits = difficulty_bits;
    int i = 0;
    while (bits > 0) {
        unsigned char byte = hash[i];
        if (bits >= 8) {
            if (byte != 0) return 0;
            bits -= 8;
        } else {
            unsigned char mask = 0xFF << (8 - bits);
            if ((byte & mask) != 0) return 0;
            bits = 0;
        }
        i++;
    }
    return 1;
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <challenge_string> <difficulty_bits>\n", argv[0]);
        return 1;
    }

    const char *challenge = argv[1];
    int difficulty_bits = atoi(argv[2]);
    if (difficulty_bits <= 0 || difficulty_bits > 256) {
        fprintf(stderr, "Difficulty bits must be between 1 and 256\n");
        return 1;
    }

    unsigned char hash[SHA256_DIGEST_LENGTH];
    unsigned long long counter = 0;

    size_t challenge_len = strlen(challenge);
    char *input = malloc(challenge_len + 32);

    if (!input) {
        fprintf(stderr, "Memory allocation failed\n");
        return 1;
    }

    while (1) {
        sprintf(input, "%s%llu", challenge, counter);
        SHA256((unsigned char *)input, strlen(input), hash);

        if (check_difficulty(hash, difficulty_bits)) {
            printf("===================\n");
            printf("Found nonce: %llu\n", counter);
            printf("===================\n");
            printf("Hash: ");
            for (int i = 0; i < SHA256_DIGEST_LENGTH; i++)
                printf("%02x", hash[i]);
            printf("\n");
            break;
        }

        if ((counter & 0xFFFFF) == 0) {
            printf("Attempts: %llu\n", counter);
        }
        counter++;
    }

    free(input);
    return 0;
}