#include "../common/des.h"

#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <mutex>
#include <string>
#include <thread>
#include <vector>

#ifdef _WIN32
#include <windows.h>
#else
#include <unistd.h>
#endif

static void fail(const char *msg) {
    std::fprintf(stderr, "%s\n", msg);
    std::exit(2);
}

static uint64_t parse_hex_u64(const char *hex) {
    const char *p = hex;
    if (p[0] == '0' && (p[1] == 'x' || p[1] == 'X')) {
        p += 2;
    }
    return std::strtoull(p, nullptr, 16);
}

static std::map<std::string, std::string> parse_opts(int argc, char **argv) {
    std::map<std::string, std::string> opts;
    for (int i = 2; i + 1 < argc; i += 2) {
        if (std::strncmp(argv[i], "--", 2) != 0) {
            i -= 1;
            continue;
        }
        opts[argv[i] + 2] = argv[i + 1];
    }
    return opts;
}

static unsigned cpu_count() {
    unsigned n = std::thread::hardware_concurrency();
    return n ? n : 1;
}

static void run_devices() {
    unsigned n = cpu_count();
    std::printf("[{\"id\":0,\"name\":\"CPU (%u cores)\",\"type\":\"cpu\",\"cores\":%u}]\n", n, n);
}

static void run_selftest() {
    if (!des_selftest()) {
        fail("DES self-test failed");
    }
    std::fprintf(stderr, "cpu_selftest ok  cores=%u\n", cpu_count());
}

static void sleep_ms(int ms) {
#ifdef _WIN32
    Sleep((DWORD)ms);
#else
    usleep((useconds_t)ms * 1000);
#endif
}

static void run_brute(const std::map<std::string, std::string> &opts) {
    uint32_t key_len = (uint32_t)std::strtoul(opts.count("key-len") ? opts.at("key-len").c_str() : "0", nullptr, 10);
    std::string charset = opts.count("charset") ? opts.at("charset") : "";
    uint32_t pad_byte = (uint32_t)(parse_hex_u64(opts.count("pad") ? opts.at("pad").c_str() : "00") & 0xFF);
    uint64_t target = parse_hex_u64(opts.count("target") ? opts.at("target").c_str() : "0");
    uint64_t start = std::strtoull(opts.count("start") ? opts.at("start").c_str() : "0", nullptr, 10);
    uint64_t count = std::strtoull(opts.count("count") ? opts.at("count").c_str() : "0", nullptr, 10);
    unsigned workers = (unsigned)std::strtoul(opts.count("workers") ? opts.at("workers").c_str() : "0", nullptr, 10);
    if (key_len == 0 || charset.empty() || count == 0) {
        fail("brute requires --key-len --charset --count");
    }
    if (workers == 0) {
        workers = cpu_count();
    }

    uint64_t fills[8] = {0};
    uint32_t nfills = 0;
    std::string fills_str = opts.count("fills") ? opts.at("fills") : "0000000000000000";
    std::string token;
    for (size_t i = 0, begin = 0; i <= fills_str.size(); i++) {
        if (i == fills_str.size() || fills_str[i] == ',') {
            token = fills_str.substr(begin, i - begin);
            if (!token.empty() && nfills < 8) {
                fills[nfills++] = parse_hex_u64(token.c_str());
            }
            begin = i + 1;
        }
    }

    const uint8_t *cs = reinterpret_cast<const uint8_t *>(charset.data());
    uint32_t clen = (uint32_t)charset.size();
    std::atomic<uint64_t> done{0};
    std::mutex out_mu;

    auto worker_fn = [&](uint64_t wstart, uint64_t wcount) {
        uint64_t local = 0;
        for (uint64_t i = 0; i < wcount; i++) {
            uint64_t index = wstart + i;
            uint64_t key = make_key(index, key_len, clen, pad_byte, cs);
            uint64_t sk[16];
            des_key_schedule(key, sk);
            uint64_t pt = des_crypt(target, sk, 1);
            if (is_fill(pt, fills, nfills)) {
                std::lock_guard<std::mutex> lock(out_mu);
                std::printf("{\"index\":%llu,\"key_hex\":\"%016llx\",\"plain_hex\":\"%016llx\"}\n",
                            (unsigned long long)index, (unsigned long long)key, (unsigned long long)pt);
                std::fflush(stdout);
            }
            local++;
            if ((local & 0x3FFF) == 0) {
                done.fetch_add(0x4000, std::memory_order_relaxed);
                local = 0;
            }
        }
        if (local) {
            done.fetch_add(local, std::memory_order_relaxed);
        }
    };

    std::vector<std::thread> threads;
    uint64_t chunk = count / workers;
    uint64_t rem = count % workers;
    uint64_t cursor = start;
    auto t0 = std::chrono::steady_clock::now();
    std::atomic<bool> finished{false};

    std::thread reporter([&]() {
        while (!finished.load(std::memory_order_relaxed)) {
            uint64_t d = done.load(std::memory_order_relaxed);
            if (d > count) {
                d = count;
            }
            auto now = std::chrono::steady_clock::now();
            double elapsed = std::chrono::duration<double>(now - t0).count();
            if (elapsed < 1e-6) {
                elapsed = 1e-6;
            }
            std::fprintf(stderr, "gpu %llu/%llu  %.0f keys/s\n", (unsigned long long)d,
                         (unsigned long long)count, d / elapsed);
            sleep_ms(200);
        }
    });

    for (unsigned w = 0; w < workers; w++) {
        uint64_t n = chunk + (w < rem ? 1 : 0);
        if (n == 0) {
            continue;
        }
        threads.emplace_back(worker_fn, cursor, n);
        cursor += n;
    }
    for (auto &t : threads) {
        t.join();
    }
    finished.store(true, std::memory_order_relaxed);
    reporter.join();

    auto now = std::chrono::steady_clock::now();
    double elapsed = std::chrono::duration<double>(now - t0).count();
    if (elapsed < 1e-6) {
        elapsed = 1e-6;
    }
    std::fprintf(stderr, "gpu %llu/%llu  %.0f keys/s\n", (unsigned long long)count,
                 (unsigned long long)count, count / elapsed);
}

int main(int argc, char **argv) {
    if (argc < 2) {
        fail("usage: med_cpu devices | selftest | brute [options]");
    }
    if (std::strcmp(argv[1], "devices") == 0) {
        run_devices();
        return 0;
    }
    auto opts = parse_opts(argc, argv);
    if (std::strcmp(argv[1], "selftest") == 0) {
        run_selftest();
        return 0;
    }
    if (std::strcmp(argv[1], "brute") == 0) {
        run_selftest();
        run_brute(opts);
        return 0;
    }
    fail("unknown command");
    return 2;
}