use sha2::{Digest, Sha256};
use std::env;
use std::process;

fn check_difficulty(hash: &[u8], difficulty_bits: u32) -> bool {
    let mut bits = difficulty_bits;
    let mut i = 0;

    while bits > 0 {
        let byte = hash[i];
        if bits >= 8 {
            if byte != 0 {
                return false;
            }
            bits -= 8;
        } else {
            let mask = 0xFF << (8 - bits);
            if (byte & mask) != 0 {
                return false;
            }
            bits = 0;
        }
        i += 1;
    }
    true
}

fn main() {
    let args: Vec<String> = env::args().collect();

    if args.len() != 3 {
        eprintln!("Usage: {} <challenge_string> <difficulty_bits>", args[0]);
        process::exit(1);
    }

    let challenge = &args[1];
    let difficulty_bits: u32 = match args[2].parse() {
        Ok(num) if num > 0 && num <= 256 => num,
        _ => {
            eprintln!("Difficulty bits must be between 1 and 256");
            process::exit(1);
        }
    };

    let mut counter: u64 = 0;

    loop {
        let input = format!("{}{}", challenge, counter);
        let hash = Sha256::digest(input.as_bytes());

        if check_difficulty(&hash, difficulty_bits) {
            println!("===================");
            println!("Found nonce: {}", counter);
            println!("===================");
            println!("Hash: {:x}", hash);
            break;
        }

        if counter & 0xFFFFF == 0 {
            println!("Attempts: {}", counter);
        }

        counter += 1;
    }
}