#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <string>
#include <vector>
#include <map>
#include <chrono>
#include <cuda_runtime.h>

enum { kMaxHits = 256, kThreadsPerGroup = 256 };

#define CUDA_CHECK(expr)                                                       \
    do {                                                                       \
        cudaError_t _err = (expr);                                             \
        if (_err != cudaSuccess) {                                             \
            std::fprintf(stderr, "CUDA error %s:%d: %s\n", __FILE__, __LINE__, \
                         cudaGetErrorString(_err));                            \
            std::exit(2);                                                      \
        }                                                                      \
    } while (0)

struct Params {
    uint64_t start;
    uint32_t key_len;
    uint32_t charset_len;
    uint32_t pad_byte;
    uint32_t fill_count;
    uint32_t batch_count;
    uint32_t _pad;
    uint64_t target;
    uint64_t fills[8];
};

struct Hit {
    uint64_t index;
    uint64_t key;
    uint64_t plain;
};

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 void json_escape(const char *src, char *dst, size_t dstsz) {
    size_t o = 0;
    for (size_t i = 0; src[i] && o + 2 < dstsz; i++) {
        unsigned char c = (unsigned char)src[i];
        if (c == '"' || c == '\\') {
            dst[o++] = '\\';
            dst[o++] = (char)c;
        } else if (c < 32) {
            dst[o++] = '?';
        } else {
            dst[o++] = (char)c;
        }
    }
    dst[o] = 0;
}

__device__ __constant__ unsigned char IP_TBL[64] = {
    58,50,42,34,26,18,10, 2, 60,52,44,36,28,20,12, 4,
    62,54,46,38,30,22,14, 6, 64,56,48,40,32,24,16, 8,
    57,49,41,33,25,17, 9, 1, 59,51,43,35,27,19,11, 3,
    61,53,45,37,29,21,13, 5, 63,55,47,39,31,23,15, 7
};
__device__ __constant__ unsigned char FP_TBL[64] = {
    40, 8,48,16,56,24,64,32, 39, 7,47,15,55,23,63,31,
    38, 6,46,14,54,22,62,30, 37, 5,45,13,53,21,61,29,
    36, 4,44,12,52,20,60,28, 35, 3,43,11,51,19,59,27,
    34, 2,42,10,50,18,58,26, 33, 1,41, 9,49,17,57,25
};
__device__ __constant__ unsigned char E_TBL[48] = {
    32, 1, 2, 3, 4, 5,  4, 5, 6, 7, 8, 9,
     8, 9,10,11,12,13, 12,13,14,15,16,17,
    16,17,18,19,20,21, 20,21,22,23,24,25,
    24,25,26,27,28,29, 28,29,30,31,32, 1
};
__device__ __constant__ unsigned char P_TBL[32] = {
    16, 7,20,21,29,12,28,17,  1,15,23,26, 5,18,31,10,
     2, 8,24,14,32,27, 3, 9, 19,13,30, 6,22,11, 4,25
};
__device__ __constant__ unsigned char PC1_TBL[56] = {
    57,49,41,33,25,17, 9,  1,58,50,42,34,26,18,
    10, 2,59,51,43,35,27, 19,11, 3,60,52,44,36,
    63,55,47,39,31,23,15,  7,62,54,46,38,30,22,
    14, 6,61,53,45,37,29, 21,13, 5,28,20,12, 4
};
__device__ __constant__ unsigned char PC2_TBL[48] = {
    14,17,11,24, 1, 5,  3,28,15, 6,21,10,
    23,19,12, 4,26, 8, 16, 7,27,20,13, 2,
    41,52,31,37,47,55, 30,40,51,45,33,48,
    44,49,39,56,34,53, 46,42,50,36,29,32
};
__device__ __constant__ unsigned char SHIFTS[16] = {1,1,2,2,2,2,2,2,1,2,2,2,2,2,2,1};
__device__ __constant__ unsigned char SBOX[8][64] = {
    {14,4,13,1,2,15,11,8,3,10,6,12,5,9,0,7, 0,15,7,4,14,2,13,1,10,6,12,11,9,5,3,8, 4,1,14,8,13,6,2,11,15,12,9,7,3,10,5,0, 15,12,8,2,4,9,1,7,5,11,3,14,10,0,6,13},
    {15,1,8,14,6,11,3,4,9,7,2,13,12,0,5,10, 3,13,4,7,15,2,8,14,12,0,1,10,6,9,11,5, 0,14,7,11,10,4,13,1,5,8,12,6,9,3,2,15, 13,8,10,1,3,15,4,2,11,6,7,12,0,5,14,9},
    {10,0,9,14,6,3,15,5,1,13,12,7,11,4,2,8, 13,7,0,9,3,4,6,10,2,8,5,14,12,11,15,1, 13,6,4,9,8,15,3,0,11,1,2,12,5,10,14,7, 1,10,13,0,6,9,8,7,4,15,14,3,11,5,2,12},
    {7,13,14,3,0,6,9,10,1,2,8,5,11,12,4,15, 13,8,11,5,6,15,0,3,4,7,2,12,1,10,14,9, 10,6,9,0,12,11,7,13,15,1,3,14,5,2,8,4, 3,15,0,6,10,1,13,8,9,4,5,11,12,7,2,14},
    {2,12,4,1,7,10,11,6,8,5,3,15,13,0,14,9, 14,11,2,12,4,7,13,1,5,0,15,10,3,9,8,6, 4,2,1,11,10,13,7,8,15,9,12,5,6,3,0,14, 11,8,12,7,1,14,2,13,6,15,0,9,10,4,5,3},
    {12,1,10,15,9,2,6,8,0,13,3,4,14,7,5,11, 10,15,4,2,7,12,9,5,6,1,13,14,0,11,3,8, 9,14,15,5,2,8,12,3,7,0,4,10,1,13,11,6, 4,3,2,12,9,5,15,10,11,14,1,7,6,0,8,13},
    {4,11,2,14,15,0,8,13,3,12,9,7,5,10,6,1, 13,0,11,7,4,9,1,10,14,3,5,12,2,15,8,6, 1,4,11,13,12,3,7,14,10,15,6,8,0,5,9,2, 6,11,13,8,1,4,10,7,9,5,0,15,14,2,3,12},
    {13,2,8,4,6,15,11,1,10,9,3,14,5,0,12,7, 1,15,13,8,10,3,7,4,12,5,6,11,0,14,9,2, 7,11,4,1,9,12,14,2,0,6,10,13,15,3,5,8, 2,1,14,7,4,10,8,13,15,12,9,0,3,5,6,11}
};

__device__ uint64_t perm(uint64_t src, const unsigned char *tbl, unsigned nout, unsigned srcbits) {
    uint64_t outv = 0;
    for (unsigned i = 0; i < nout; i++) {
        uint64_t bit = (src >> (srcbits - tbl[i])) & 1ULL;
        outv = (outv << 1) | bit;
    }
    return outv;
}

__device__ unsigned rotl28(unsigned v, unsigned s) {
    return ((v << s) | (v >> (28 - s))) & 0x0FFFFFFFu;
}

__device__ void des_key_schedule(uint64_t key, uint64_t sk[16]) {
    uint64_t cd = perm(key, PC1_TBL, 56, 64);
    unsigned c = (unsigned)(cd >> 28);
    unsigned d = (unsigned)(cd & 0x0FFFFFFFULL);
    for (unsigned r = 0; r < 16; r++) {
        c = rotl28(c, SHIFTS[r]);
        d = rotl28(d, SHIFTS[r]);
        uint64_t cd2 = ((uint64_t)c << 28) | (uint64_t)d;
        sk[r] = perm(cd2, PC2_TBL, 48, 56);
    }
}

__device__ unsigned feistel(unsigned r, uint64_t subkey) {
    uint64_t er = perm((uint64_t)r, E_TBL, 48, 32) ^ subkey;
    unsigned s = 0;
    for (unsigned i = 0; i < 8; i++) {
        unsigned chunk = (unsigned)((er >> (42 - 6 * i)) & 0x3F);
        unsigned row = ((chunk & 0x20) >> 4) | (chunk & 1);
        unsigned col = (chunk >> 1) & 0xF;
        s = (s << 4) | (unsigned)SBOX[i][row * 16 + col];
    }
    return (unsigned)perm((uint64_t)s, P_TBL, 32, 32);
}

__device__ uint64_t des_crypt(uint64_t block, const uint64_t sk[16], int decrypt) {
    uint64_t ip = perm(block, IP_TBL, 64, 64);
    unsigned l = (unsigned)(ip >> 32);
    unsigned r = (unsigned)(ip & 0xFFFFFFFFULL);
    for (unsigned i = 0; i < 16; i++) {
        unsigned rnd = decrypt ? (15 - i) : i;
        unsigned n = l ^ feistel(r, sk[rnd]);
        l = r;
        r = n;
    }
    uint64_t pre = ((uint64_t)r << 32) | (uint64_t)l;
    return perm(pre, FP_TBL, 64, 64);
}

__device__ uint64_t make_key(uint64_t index, const Params &p, const unsigned char *charset) {
    unsigned char bytes[8];
    for (int i = 0; i < 8; i++) {
        bytes[i] = (unsigned char)p.pad_byte;
    }
    uint64_t n = index;
    unsigned clen = p.charset_len;
    for (int pos = (int)p.key_len - 1; pos >= 0; pos--) {
        bytes[pos] = charset[n % clen];
        n /= clen;
    }
    uint64_t key = 0;
    for (int i = 0; i < 8; i++) {
        key = (key << 8) | (uint64_t)bytes[i];
    }
    return key;
}

__device__ int is_fill(uint64_t pt, const Params &p) {
    for (unsigned i = 0; i < p.fill_count; i++) {
        if (pt == p.fills[i]) {
            return 1;
        }
    }
    return 0;
}

__global__ void des_known_answer(uint64_t *out) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;
    uint64_t key = 0x133457799BBCDFF1ULL;
    uint64_t pt = 0x0123456789ABCDEFULL;
    uint64_t sk[16];
    des_key_schedule(key, sk);
    out[0] = des_crypt(pt, sk, 0);
    out[1] = des_crypt(out[0], sk, 1);
}

__global__ void des_brute(Params params, const unsigned char *charset, Hit *hits, unsigned *hit_count) {
    unsigned gid = blockIdx.x * blockDim.x + threadIdx.x;
    if (gid >= params.batch_count) return;
    uint64_t index = params.start + (uint64_t)gid;
    uint64_t key = make_key(index, params, charset);
    uint64_t sk[16];
    des_key_schedule(key, sk);
    uint64_t pt = des_crypt(params.target, sk, 1);
    if (!is_fill(pt, params)) return;
    unsigned slot = atomicAdd(hit_count, 1u);
    if (slot < 256) {
        hits[slot].index = index;
        hits[slot].key = key;
        hits[slot].plain = pt;
    }
}

static int pick_device(const std::map<std::string, std::string> &opts) {
    int count = 0;
    CUDA_CHECK(cudaGetDeviceCount(&count));
    if (count <= 0) fail("no CUDA GPU");
    int idx = opts.count("device") ? std::atoi(opts.at("device").c_str()) : 0;
    if (idx < 0 || idx >= count) idx = 0;
    CUDA_CHECK(cudaSetDevice(idx));
    return idx;
}

static void run_devices() {
    int count = 0;
    cudaError_t err = cudaGetDeviceCount(&count);
    if (err != cudaSuccess) {
        std::printf("[]\n");
        return;
    }
    std::printf("[");
    for (int i = 0; i < count; i++) {
        cudaDeviceProp prop;
        CUDA_CHECK(cudaGetDeviceProperties(&prop, i));
        char name[512];
        json_escape(prop.name, name, sizeof(name));
        std::printf("%s{\"id\":%d,\"name\":\"%s\",\"type\":\"cuda\",\"memory\":%llu}",
                    i ? "," : "", i, name, (unsigned long long)prop.totalGlobalMem);
    }
    std::printf("]\n");
}

static void run_selftest() {
    uint64_t *out = nullptr;
    CUDA_CHECK(cudaMalloc(&out, 16));
    des_known_answer<<<1, 1>>>(out);
    CUDA_CHECK(cudaDeviceSynchronize());
    uint64_t words[2];
    CUDA_CHECK(cudaMemcpy(words, out, 16, cudaMemcpyDeviceToHost));
    CUDA_CHECK(cudaFree(out));
    uint64_t expected = 0x85E813540F0AB405ULL;
    uint64_t plain = 0x0123456789ABCDEFULL;
    if (words[0] != expected || words[1] != plain) {
        std::fprintf(stderr, "DES self-test failed ct=%016llx expected=%016llx roundtrip=%016llx\n",
                     (unsigned long long)words[0], (unsigned long long)expected,
                     (unsigned long long)words[1]);
        std::exit(2);
    }
    cudaDeviceProp prop;
    int dev = 0;
    CUDA_CHECK(cudaGetDevice(&dev));
    CUDA_CHECK(cudaGetDeviceProperties(&prop, dev));
    std::fprintf(stderr, "cuda_selftest ok  gpu=%s\n", prop.name);
}

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);
    uint64_t batch = std::strtoull(opts.count("batch") ? opts.at("batch").c_str() : "16777216", nullptr, 10);
    if (key_len == 0 || charset.empty() || count == 0) {
        fail("brute requires --key-len --charset --count");
    }

    Params params{};
    params.key_len = key_len;
    params.charset_len = (uint32_t)charset.size();
    params.pad_byte = pad_byte;
    params.target = target;
    std::string fills_str = opts.count("fills") ? opts.at("fills") : "0000000000000000";
    uint32_t nfills = 0;
    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) {
                params.fills[nfills++] = parse_hex_u64(token.c_str());
            }
            begin = i + 1;
        }
    }
    params.fill_count = nfills;

    unsigned char *d_charset = nullptr;
    Hit *d_hits = nullptr;
    unsigned *d_hitn = nullptr;
    CUDA_CHECK(cudaMalloc(&d_charset, charset.size()));
    CUDA_CHECK(cudaMemcpy(d_charset, charset.data(), charset.size(), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMalloc(&d_hits, kMaxHits * sizeof(Hit)));
    CUDA_CHECK(cudaMalloc(&d_hitn, sizeof(unsigned)));

    uint64_t done = 0;
    auto t0 = std::chrono::steady_clock::now();
    std::vector<Hit> host_hits(kMaxHits);
    while (done < count) {
        uint64_t n = batch < (count - done) ? batch : (count - done);
        params.start = start + done;
        params.batch_count = (uint32_t)n;
        unsigned zero = 0;
        CUDA_CHECK(cudaMemcpy(d_hitn, &zero, sizeof(zero), cudaMemcpyHostToDevice));
        unsigned groups = (unsigned)((n + kThreadsPerGroup - 1) / kThreadsPerGroup);
        des_brute<<<groups, kThreadsPerGroup>>>(params, d_charset, d_hits, d_hitn);
        CUDA_CHECK(cudaDeviceSynchronize());
        unsigned hitn = 0;
        CUDA_CHECK(cudaMemcpy(&hitn, d_hitn, sizeof(hitn), cudaMemcpyDeviceToHost));
        if (hitn > kMaxHits) hitn = kMaxHits;
        if (hitn) {
            CUDA_CHECK(cudaMemcpy(host_hits.data(), d_hits, hitn * sizeof(Hit), cudaMemcpyDeviceToHost));
            for (unsigned i = 0; i < hitn; i++) {
                std::printf("{\"index\":%llu,\"key_hex\":\"%016llx\",\"plain_hex\":\"%016llx\"}\n",
                            (unsigned long long)host_hits[i].index,
                            (unsigned long long)host_hits[i].key,
                            (unsigned long long)host_hits[i].plain);
                std::fflush(stdout);
            }
        }
        done += n;
        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)done,
                     (unsigned long long)count, done / elapsed);
    }

    CUDA_CHECK(cudaFree(d_charset));
    CUDA_CHECK(cudaFree(d_hits));
    CUDA_CHECK(cudaFree(d_hitn));
}

int main(int argc, char **argv) {
    if (argc < 2) {
        fail("usage: med_cuda devices | selftest [--device N] | brute [--device N] [options]");
    }
    if (std::strcmp(argv[1], "devices") == 0) {
        run_devices();
        return 0;
    }
    auto opts = parse_opts(argc, argv);
    pick_device(opts);
    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;
}