#include <randomx.h>

#include <algorithm>
#include <array>
#include <charconv>
#include <cstdio>
#include <cstdint>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

#ifdef _WIN32
#include <fcntl.h>
#include <io.h>
#endif

namespace {
constexpr std::size_t MaxBlob = 8192;
constexpr std::size_t MaxLine = 16600;
using Hash = std::array<std::uint8_t, RANDOMX_HASH_SIZE>;

[[noreturn]] void invalid() { throw std::runtime_error("invalid command"); }

int nibble(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;
    invalid();
}

std::vector<std::uint8_t> unhex(std::string_view text, std::size_t min, std::size_t max) {
    if (text.size() % 2 != 0 || text.size() / 2 < min || text.size() / 2 > max) invalid();
    std::vector<std::uint8_t> bytes(text.size() / 2);
    for (std::size_t i = 0; i < bytes.size(); ++i)
        bytes[i] = static_cast<std::uint8_t>((nibble(text[2 * i]) << 4) | nibble(text[2 * i + 1]));
    return bytes;
}

std::string hex(const Hash& hash) {
    constexpr char digits[] = "0123456789abcdef";
    std::string out(hash.size() * 2, '0');
    for (std::size_t i = 0; i < hash.size(); ++i) {
        out[2 * i] = digits[hash[i] >> 4];
        out[2 * i + 1] = digits[hash[i] & 15];
    }
    return out;
}

std::uint64_t decimal(std::string_view text) {
    if (text.empty() || text.size() > 20) invalid();
    for (char c : text) if (c < '0' || c > '9') invalid();
    std::uint64_t value = 0;
    const auto result = std::from_chars(text.data(), text.data() + text.size(), value);
    if (result.ec != std::errc{} || result.ptr != text.data() + text.size()) invalid();
    return value;
}

// Read incrementally into bounded storage; getline(string) alone allows unbounded allocation.
bool readLine(std::string& line) {
    line.clear();
    char c;
    while (std::cin.get(c)) {
        if (c == '\n') return true;
        if (line.size() >= MaxLine || static_cast<unsigned char>(c) > 127 || c == '\0') invalid();
        line.push_back(c);
    }
    if (!std::cin.eof()) throw std::runtime_error("stdin failure");
    if (!line.empty()) throw std::runtime_error("truncated command");
    return false;
}

std::vector<std::string_view> tokens(std::string& line) {
    if (!line.empty() && line.back() == '\r') line.pop_back();
    std::vector<std::string_view> parts;
    std::string_view remaining(line);
    while (true) {
        const auto space = remaining.find(' ');
        const auto part = remaining.substr(0, space);
        if (part.empty() || parts.size() == 6) invalid();
        parts.push_back(part);
        if (space == std::string_view::npos) break;
        remaining.remove_prefix(space + 1);
    }
    return parts;
}

class Engine {
    using Cache = std::unique_ptr<randomx_cache, decltype(&randomx_release_cache)>;
    using Dataset = std::unique_ptr<randomx_dataset, decltype(&randomx_release_dataset)>;
    using Vm = std::unique_ptr<randomx_vm, decltype(&randomx_destroy_vm)>;
    bool fast_;
    randomx_flags flags_ = randomx_get_flags() | RANDOMX_FLAG_SECURE;
    // Destruction order: VM, dataset, cache. No borrowed pointer outlives its allocation.
    Cache cache_{nullptr, randomx_release_cache};
    Dataset dataset_{nullptr, randomx_release_dataset};
    Vm vm_{nullptr, randomx_destroy_vm};
    std::vector<std::uint8_t> key_;

public:
    explicit Engine(bool fast) : fast_(fast) {}

    void key(const std::vector<std::uint8_t>& key) {
        if (vm_ && key == key_) return;
        if (!cache_) {
            cache_.reset(randomx_alloc_cache(flags_));
            if (!cache_) throw std::runtime_error("RandomX cache allocation failed");
        }
        // Destroy the old VM before changing memory it references; reuse allocated cache/dataset.
        vm_.reset();
        randomx_init_cache(cache_.get(), key.data(), key.size());
        if (fast_) {
            if (!dataset_) dataset_.reset(randomx_alloc_dataset(RANDOMX_FLAG_DEFAULT));
            if (!dataset_) throw std::runtime_error("RandomX dataset allocation failed");
            randomx_init_dataset(dataset_.get(), cache_.get(), 0, randomx_dataset_item_count());
        }
        const auto vmFlags = fast_ ? flags_ | RANDOMX_FLAG_FULL_MEM : flags_;
        vm_.reset(randomx_create_vm(vmFlags, fast_ ? nullptr : cache_.get(), dataset_.get()));
        if (!vm_) throw std::runtime_error("RandomX VM allocation failed");
        key_ = key;
    }

    Hash hash(const std::vector<std::uint8_t>& blob) {
        Hash result{};
        randomx_calculate_hash(vm_.get(), blob.data(), blob.size(), result.data());
        return result;
    }
};

void command(Engine& engine, const std::vector<std::string_view>& parts) {
    if (parts.size() != 3 && parts.size() != 6) invalid();
    const bool mining = parts[0] == "mine";
    if ((!mining && parts[0] != "hash") || parts.size() != (mining ? 6U : 3U)) invalid();
    const auto key = unhex(parts[1], 32, 32);
    auto blob = unhex(parts[2], 1, MaxBlob);
    // All validation precedes expensive cache/dataset initialization.
    if (!mining) {
        engine.key(key);
        std::cout << hex(engine.hash(blob)) << std::endl;
        return;
    }
    const auto target = unhex(parts[3], 32, 32);
    const auto offsetValue = decimal(parts[4]);
    const auto attempts = decimal(parts[5]);
    if (blob.size() < 8 || offsetValue > blob.size() - 8 || attempts < 1 || attempts > 10000) invalid();
    const auto offset = static_cast<std::size_t>(offsetValue);
    std::uint64_t nonce = 0;
    for (unsigned i = 0; i < 8; ++i) nonce |= std::uint64_t{blob[offset + i]} << (8 * i);
    engine.key(key);
    for (std::uint64_t attempt = 0; attempt < attempts; ++attempt) {
        for (unsigned i = 0; i < 8; ++i)
            blob[offset + i] = static_cast<std::uint8_t>(nonce >> (8 * i));
        const auto hash = engine.hash(blob);
        // BIG ENDIAN integer comparison (intentionally NOT Monero's endian convention).
        if (!std::lexicographical_compare(target.begin(), target.end(), hash.begin(), hash.end())) {
            std::cout << "found " << hex(hash) << ' ' << nonce << std::endl;
            return;
        }
        if (nonce == std::numeric_limits<std::uint64_t>::max()) break; // Never wrap/repeat work.
        ++nonce;
    }
    std::cout << "exhausted" << std::endl;
}

void selfTest(Engine& engine) {
    // Exact upstream v1.2.1 Hash test 1a/2a, including the original 12-byte key.
    // This argv-only test does not weaken the protocol's fixed 32-byte key validation.
    const std::string key = "test key 000";
    const std::string blob = "This is a test";
    engine.key(std::vector<std::uint8_t>(key.begin(), key.end()));
    if (hex(engine.hash(std::vector<std::uint8_t>(blob.begin(), blob.end()))) !=
        "639183aae1bf4c9a35884cb46b09cad9175f04efd7684e7262a0ac1c2f0b4e3f")
        throw std::runtime_error("upstream vector mismatch");
    std::cerr << "RandomX upstream vector passed\n";
}
} // namespace

int main(int argc, char** argv) {
    try {
#ifdef _WIN32
        // CRT text-mode Ctrl-Z must not silently truncate an untrusted command stream.
        if (_setmode(_fileno(stdin), _O_BINARY) == -1 ||
            _setmode(_fileno(stdout), _O_BINARY) == -1)
            throw std::runtime_error("stdio binary mode failed");
#endif
        std::ios::sync_with_stdio(false);
        std::cin.tie(nullptr);
        bool fast = false;
        bool test = false;
        for (int i = 1; i < argc; ++i) {
            const std::string_view arg(argv[i]);
            if (arg == "--fast" && !fast) fast = true;
            else if (arg == "--self-test" && !test) test = true;
            else throw std::runtime_error("usage: wikideck-randomx [--fast] [--self-test]");
        }
        Engine engine(fast);
        if (test) { selfTest(engine); return 0; }
        std::string line;
        while (readLine(line)) {
            command(engine, tokens(line));
            if (!std::cout) throw std::runtime_error("stdout failure");
        }
        return 0;
    } catch (const std::exception& error) {
        std::cerr << "randomx: " << error.what() << '\n';
        return 1;
    } catch (...) {
        std::cerr << "randomx: native failure\n";
        return 1;
    }
}