#include <openssl/sha.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <exception>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <limits>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

namespace {

using Header = std::array<unsigned char, 80>;
using Hash = std::array<unsigned char, 32>;

int HexDigit(char c) {
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return -1;
}

template <std::size_t N>
std::array<unsigned char, N> ParseHex(const std::string& text,
                                     const char* name) {
    if (text.size() != N * 2) {
        throw std::invalid_argument(std::string(name) + " must contain exactly " +
                                    std::to_string(N * 2) + " hexadecimal characters");
    }
    std::array<unsigned char, N> value{};
    for (std::size_t i = 0; i < N; ++i) {
        const int hi = HexDigit(text[2 * i]);
        const int lo = HexDigit(text[2 * i + 1]);
        if (hi < 0 || lo < 0) {
            throw std::invalid_argument(std::string(name) + " contains a non-hexadecimal character");
        }
        value[i] = static_cast<unsigned char>((hi << 4) | lo);
    }
    return value;
}

template <std::size_t N>
std::string Hex(const std::array<unsigned char, N>& value, bool reverse = false) {
    static constexpr char digits[] = "0123456789abcdef";
    std::string result(N * 2, '0');
    for (std::size_t i = 0; i < N; ++i) {
        const unsigned char byte = value[reverse ? N - 1 - i : i];
        result[2 * i] = digits[byte >> 4];
        result[2 * i + 1] = digits[byte & 15];
    }
    return result;
}

std::uint64_t ParseUnsigned(const std::string& text, const char* name,
                            std::uint64_t minimum, std::uint64_t maximum) {
    if (text.empty()) throw std::invalid_argument(std::string(name) + " is empty");
    std::uint64_t value = 0;
    for (char c : text) {
        if (c < '0' || c > '9') {
            throw std::invalid_argument(std::string(name) + " must be an unsigned decimal integer");
        }
        const unsigned digit = static_cast<unsigned>(c - '0');
        if (value > (std::numeric_limits<std::uint64_t>::max() - digit) / 10) {
            throw std::invalid_argument(std::string(name) + " is too large");
        }
        value = value * 10 + digit;
    }
    if (value < minimum || value > maximum) {
        throw std::invalid_argument(std::string(name) + " must be between " +
                                    std::to_string(minimum) + " and " +
                                    std::to_string(maximum));
    }
    return value;
}

void SetNonce(unsigned char* bytes, std::uint32_t nonce) {
    bytes[0] = static_cast<unsigned char>(nonce);
    bytes[1] = static_cast<unsigned char>(nonce >> 8);
    bytes[2] = static_cast<unsigned char>(nonce >> 16);
    bytes[3] = static_cast<unsigned char>(nonce >> 24);
}

// SHA256_Transform consumes and updates h only; padding and message lengths are
// supplied explicitly below. The shared first 64 header bytes need hashing once.
struct HeaderHasher {
    SHA256_CTX midstate{};
    SHA256_CTX initial{};
    std::array<unsigned char, 64> tail{};

    explicit HeaderHasher(const Header& header) {
        SHA256_Init(&initial);
        midstate = initial;
        SHA256_Transform(&midstate, header.data());
        std::copy(header.begin() + 64, header.end(), tail.begin());
        tail[16] = 0x80;
        tail[62] = 0x02;  // 80 bytes * 8 = 640 bits, big-endian.
        tail[63] = 0x80;
    }
};

void SerializeState(const SHA256_CTX& context, unsigned char* output) {
    for (unsigned i = 0; i < 8; ++i) {
        const std::uint32_t word = context.h[i];
        output[4 * i] = static_cast<unsigned char>(word >> 24);
        output[4 * i + 1] = static_cast<unsigned char>(word >> 16);
        output[4 * i + 2] = static_cast<unsigned char>(word >> 8);
        output[4 * i + 3] = static_cast<unsigned char>(word);
    }
}

struct WorkerHasher {
    const HeaderHasher& common;
    SHA256_CTX first{};
    SHA256_CTX second{};
    std::array<unsigned char, 64> tail;
    std::array<unsigned char, 64> second_block{};

    explicit WorkerHasher(const HeaderHasher& hasher)
        : common(hasher), first(hasher.initial), second(hasher.initial), tail(hasher.tail) {
        second_block[32] = 0x80;
        second_block[62] = 0x01;  // 32 bytes * 8 = 256 bits, big-endian.
    }

    void Run() {
        std::copy(std::begin(common.midstate.h), std::end(common.midstate.h), std::begin(first.h));
        SHA256_Transform(&first, tail.data());
        SerializeState(first, second_block.data());
        std::copy(std::begin(common.initial.h), std::end(common.initial.h), std::begin(second.h));
        SHA256_Transform(&second, second_block.data());
    }

    void RunNonce(std::uint32_t nonce) {
        SetNonce(tail.data() + 12, nonce);
        Run();
    }

    Hash Digest() const {
        Hash digest{};
        SerializeState(second, digest.data());
        return digest;
    }

    bool MeetsTarget(const Hash& target) const {
        // Bitcoin treats the raw SHA digest as a little-endian uint256. Compare
        // its reversed bytes directly against the caller's big-endian target.
        for (int word_index = 7, target_index = 0; word_index >= 0; --word_index) {
            const std::uint32_t word = second.h[word_index];
            for (unsigned shift = 0; shift < 32; shift += 8, ++target_index) {
                const unsigned char byte = static_cast<unsigned char>(word >> shift);
                if (byte < target[target_index]) return true;
                if (byte > target[target_index]) return false;
            }
        }
        return true;  // Equality satisfies proof of work.
    }
};

std::string JsonEscape(const std::string& input) {
    std::string output;
    for (unsigned char c : input) {
        if (c == '"' || c == '\\') output += '\\';
        if (c < 0x20) {
            static constexpr char digits[] = "0123456789abcdef";
            output += "\\u00";
            output += digits[c >> 4];
            output += digits[c & 15];
        } else {
            output += static_cast<char>(c);
        }
    }
    return output;
}

int Solve(Header header, const Hash& target, std::uint64_t maximum_nonces,
          unsigned thread_count) {
    const HeaderHasher hasher(header);
    std::atomic<bool> found{false};
    std::uint32_t winning_nonce = 0;
    Hash winning_hash{};
    std::vector<std::uint64_t> counts(thread_count, 0);
    std::vector<std::thread> workers;
    workers.reserve(thread_count);
    const auto start = std::chrono::steady_clock::now();
    try {
        for (unsigned worker = 0; worker < thread_count; ++worker) {
            workers.emplace_back([&, worker] {
                WorkerHasher local(hasher);
                std::uint64_t attempts = 0;
                // MAX_NONCES counts candidates starting at zero. A 64-bit loop
                // index allows a full 2^32 scan without nonce wraparound.
                for (std::uint64_t nonce = worker; nonce < maximum_nonces; nonce += thread_count) {
                    if (found.load(std::memory_order_relaxed)) break;
                    local.RunNonce(static_cast<std::uint32_t>(nonce));
                    ++attempts;
                    if (local.MeetsTarget(target)) {
                        bool expected = false;
                        if (found.compare_exchange_strong(expected, true, std::memory_order_relaxed)) {
                            winning_nonce = static_cast<std::uint32_t>(nonce);
                            winning_hash = local.Digest();
                        }
                        break;
                    }
                }
                counts[worker] = attempts;
            });
        }
    } catch (...) {
        found.store(true, std::memory_order_relaxed);
        for (auto& worker : workers) worker.join();
        throw;
    }
    for (auto& worker : workers) worker.join();
    const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
    std::uint64_t hashes = 0;
    for (std::uint64_t count : counts) hashes += count;
    const bool solved = found.load(std::memory_order_relaxed);

    std::cout << std::setprecision(9) << "{\"found\":" << (solved ? "true" : "false");
    if (solved) {
        SetNonce(header.data() + 76, winning_nonce);
        std::cout << ",\"nonce\":" << winning_nonce
                  << ",\"hash\":\"" << Hex(winning_hash, true)
                  << "\",\"header\":\"" << Hex(header) << '"';
    } else {
        std::cout << ",\"nonce\":null,\"hash\":null,\"header\":null";
    }
    std::cout << ",\"hashes\":" << hashes << ",\"seconds\":" << seconds << "}\n";
    return solved ? 0 : 2;
}

}  // namespace

int main(int argc, char** argv) {
    try {
        if (argc == 3 && std::string(argv[1]) == "--hash") {
            const Header header = ParseHex<80>(argv[2], "HEADER_HEX");
            const HeaderHasher common(header);
            WorkerHasher hasher(common);
            hasher.Run();
            std::cout << "{\"hash\":\"" << Hex(hasher.Digest(), true) << "\"}\n";
            return 0;
        }
        if (argc == 6 && std::string(argv[1]) == "--solve") {
            const Header header = ParseHex<80>(argv[2], "HEADER_HEX");
            const Hash target = ParseHex<32>(argv[3], "TARGET_BIG_ENDIAN_HEX");
            if (std::all_of(target.begin(), target.end(), [](unsigned char byte) { return byte == 0; })) {
                throw std::invalid_argument("TARGET_BIG_ENDIAN_HEX must be nonzero");
            }
            const std::uint64_t maximum_nonces = ParseUnsigned(argv[4], "MAX_NONCES", 1, UINT64_C(4294967296));
            const unsigned thread_count = static_cast<unsigned>(ParseUnsigned(argv[5], "THREADS", 1, 2));
            return Solve(header, target, maximum_nonces, thread_count);
        }
        throw std::invalid_argument(
            "usage: header_grinder --hash HEADER_HEX | "
            "header_grinder --solve HEADER_HEX TARGET_BIG_ENDIAN_HEX MAX_NONCES THREADS");
    } catch (const std::exception& error) {
        std::cerr << "{\"error\":\"" << JsonEscape(error.what()) << "\"}\n";
        return 1;
    }
}