package main

import (
    "crypto/sha256"
    "encoding/hex"
    "fmt"
    "os"
    "strconv"
)

func checkDifficulty(hash []byte, difficultyBits int) bool {
    bits := difficultyBits
    i := 0
    for bits > 0 {
        b := hash[i]
        if bits >= 8 {
            if b != 0 {
                return false
            }
            bits -= 8
        } else {
            mask := byte(0xFF << (8 - bits))
            if (b & mask) != 0 {
                return false
            }
            bits = 0
        }
        i++
    }
    return true
}

func main() {
    if len(os.Args) != 3 {
        fmt.Fprintf(os.Stderr, "Usage: %s <challenge_string> <difficulty_bits>\n", os.Args[0])
        os.Exit(1)
    }

    challenge := os.Args[1]
    difficultyBits, err := strconv.Atoi(os.Args[2])
    if err != nil || difficultyBits <= 0 || difficultyBits > 256 {
        fmt.Fprintln(os.Stderr, "Difficulty bits must be between 1 and 256")
        os.Exit(1)
    }

    var counter uint64 = 0

    for {
        input := fmt.Sprintf("%s%d", challenge, counter)
        hash := sha256.Sum256([]byte(input))

        if checkDifficulty(hash[:], difficultyBits) {
            fmt.Printf("===================\n")
            fmt.Printf("Found nonce: %d\n", counter)
            fmt.Printf("===================\n")
            fmt.Printf("Hash: %s\n", hex.EncodeToString(hash[:]))
            break
        }

        if counter&0xFFFFF == 0 {
            fmt.Printf("Attempts: %d\n", counter)
        }

        counter++
    }
}