// Originally by: https://github.com/boykopovar/aiodeepseek

#include "pow.h"

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <exception>
#include <immintrin.h>
#include <string>
#include <thread>
#include <vector>

static constexpr size_t RATE = 136;
static constexpr size_t RATE_W = RATE / 8;
static constexpr size_t MAX_NONCE_DEC = 20;

/*
 * Динамическая область должна учитывать:
 *
 * base_len % 8
 * + nonce до 20 символов
 * + байт padding 0x06
 *
 * 24 байт в исходном варианте может оказаться недостаточно.
 */
static constexpr size_t DYN_SIZE = 32;

static const uint64_t RC[24] = {
    0x0000000000000001ULL,
    0x0000000000008082ULL,
    0x800000000000808AULL,
    0x8000000080008000ULL,
    0x000000000000808BULL,
    0x0000000080000001ULL,
    0x8000000080008081ULL,
    0x8000000000008009ULL,
    0x000000000000008AULL,
    0x0000000000000088ULL,
    0x0000000080008009ULL,
    0x000000008000000AULL,
    0x000000008000808BULL,
    0x800000000000008BULL,
    0x8000000000008089ULL,
    0x8000000000008003ULL,
    0x8000000000008002ULL,
    0x8000000000000080ULL,
    0x000000000000800AULL,
    0x800000008000000AULL,
    0x8000000080008081ULL,
    0x8000000000008080ULL,
    0x0000000080000001ULL,
    0x8000000080008008ULL,
};

static constexpr char DIGITS_00_99[] =
    "00010203040506070809"
    "10111213141516171819"
    "20212223242526272829"
    "30313233343536373839"
    "40414243444546474849"
    "50515253545556575859"
    "60616263646566676869"
    "70717273747576777879"
    "80818283848586878889"
    "90919293949596979899";

static inline int hex_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;
    }

    return -1;
}

static bool valid_hex_digest(const std::string& value) {
    if (value.size() != 64) {
        return false;
    }

    for (char c : value) {
        if (hex_nibble(c) < 0) {
            return false;
        }
    }

    return true;
}

static inline int fast_u64_to_dec(uint64_t value, char* out) {
    char tmp[20];
    char* p = tmp + sizeof(tmp);

    while (value >= 100) {
        const uint64_t quotient = value / 100;
        const uint64_t remainder = value - quotient * 100;

        p -= 2;
        p[0] = DIGITS_00_99[remainder * 2];
        p[1] = DIGITS_00_99[remainder * 2 + 1];

        value = quotient;
    }

    if (value < 10) {
        *--p = static_cast<char>('0' + value);
    } else {
        p -= 2;
        p[0] = DIGITS_00_99[value * 2];
        p[1] = DIGITS_00_99[value * 2 + 1];
    }

    const int len = static_cast<int>(
        tmp + sizeof(tmp) - p
    );

    std::memcpy(
        out,
        p,
        static_cast<size_t>(len)
    );

    return len;
}

struct PowCtx {
    uint64_t static_state[25];
    uint64_t target4[4];

    uint8_t dyn_tpl[DYN_SIZE];

    size_t base_len;
    size_t dyn_word_start;
    size_t dyn_word_count;
    size_t dyn_offset;
};

static PowCtx build_ctx(
    const std::string& base,
    const std::string& challenge_hex
) {
    PowCtx ctx = {};
    ctx.base_len = base.size();

    for (int word = 0; word < 4; ++word) {
        uint64_t value = 0;

        for (int byte = 0; byte < 8; ++byte) {
            const int position = (word * 8 + byte) * 2;

            const auto high = static_cast<uint8_t>(
                hex_nibble(challenge_hex[position])
            );

            const auto low = static_cast<uint8_t>(
                hex_nibble(challenge_hex[position + 1])
            );

            const uint8_t decoded = static_cast<uint8_t>(
                (high << 4) | low
            );

            value |= static_cast<uint64_t>(decoded)
                     << (byte * 8);
        }

        ctx.target4[word] = value;
    }

    ctx.dyn_word_start = ctx.base_len / 8;
    ctx.dyn_offset = ctx.base_len % 8;

    const size_t dyn_last =
        ctx.base_len +
        MAX_NONCE_DEC +
        1;

    const size_t dyn_word_end =
        dyn_last / 8 + 1;

    ctx.dyn_word_count =
        dyn_word_end - ctx.dyn_word_start;

    if (
        ctx.dyn_word_start +
        ctx.dyn_word_count >
        RATE_W
    ) {
        ctx.dyn_word_count =
            RATE_W - ctx.dyn_word_start;
    }

    std::memset(
        ctx.dyn_tpl,
        0,
        sizeof(ctx.dyn_tpl)
    );

    const size_t dyn_byte_start =
        ctx.dyn_word_start * 8;

    if (ctx.base_len > dyn_byte_start) {
        const size_t prefix_len =
            ctx.base_len - dyn_byte_start;

        std::memcpy(
            ctx.dyn_tpl,
            base.data() + dyn_byte_start,
            prefix_len
        );
    }

    alignas(8) uint8_t block[RATE] = {};

    std::memcpy(
        block,
        base.data(),
        ctx.base_len
    );

    /*
     * Последний SHA3 padding-бит.
     * Первый байт padding 0x06 добавляется после nonce.
     */
    block[RATE - 1] = 0x80;

    std::memset(
        ctx.static_state,
        0,
        sizeof(ctx.static_state)
    );

    for (size_t word = 0; word < RATE_W; ++word) {
        const bool dynamic =
            word >= ctx.dyn_word_start &&
            word <
                ctx.dyn_word_start +
                ctx.dyn_word_count;

        if (dynamic) {
            continue;
        }

        uint64_t value;

        std::memcpy(
            &value,
            &block[word * 8],
            sizeof(value)
        );

        ctx.static_state[word] ^= value;
    }

    return ctx;
}

#define R4(x, n) \
    _mm256_or_si256( \
        _mm256_slli_epi64((x), (n)), \
        _mm256_srli_epi64((x), 64 - (n)) \
    )

#define KF_ROUND4(i) \
    do { \
        C[0] = _mm256_xor_si256( \
            _mm256_xor_si256( \
                _mm256_xor_si256( \
                    _mm256_xor_si256(A[0], A[5]), \
                    A[10] \
                ), \
                A[15] \
            ), \
            A[20] \
        ); \
        C[1] = _mm256_xor_si256( \
            _mm256_xor_si256( \
                _mm256_xor_si256( \
                    _mm256_xor_si256(A[1], A[6]), \
                    A[11] \
                ), \
                A[16] \
            ), \
            A[21] \
        ); \
        C[2] = _mm256_xor_si256( \
            _mm256_xor_si256( \
                _mm256_xor_si256( \
                    _mm256_xor_si256(A[2], A[7]), \
                    A[12] \
                ), \
                A[17] \
            ), \
            A[22] \
        ); \
        C[3] = _mm256_xor_si256( \
            _mm256_xor_si256( \
                _mm256_xor_si256( \
                    _mm256_xor_si256(A[3], A[8]), \
                    A[13] \
                ), \
                A[18] \
            ), \
            A[23] \
        ); \
        C[4] = _mm256_xor_si256( \
            _mm256_xor_si256( \
                _mm256_xor_si256( \
                    _mm256_xor_si256(A[4], A[9]), \
                    A[14] \
                ), \
                A[19] \
            ), \
            A[24] \
        ); \
        \
        D[0] = _mm256_xor_si256(C[4], R4(C[1], 1)); \
        D[1] = _mm256_xor_si256(C[0], R4(C[2], 1)); \
        D[2] = _mm256_xor_si256(C[1], R4(C[3], 1)); \
        D[3] = _mm256_xor_si256(C[2], R4(C[4], 1)); \
        D[4] = _mm256_xor_si256(C[3], R4(C[0], 1)); \
        \
        for (int j = 0; j < 25; ++j) { \
            A[j] = _mm256_xor_si256(A[j], D[j % 5]); \
        } \
        \
        B[0]  = A[0]; \
        B[1]  = R4(A[6], 44); \
        B[2]  = R4(A[12], 43); \
        B[3]  = R4(A[18], 21); \
        B[4]  = R4(A[24], 14); \
        B[5]  = R4(A[3], 28); \
        B[6]  = R4(A[9], 20); \
        B[7]  = R4(A[10], 3); \
        B[8]  = R4(A[16], 45); \
        B[9]  = R4(A[22], 61); \
        B[10] = R4(A[1], 1); \
        B[11] = R4(A[7], 6); \
        B[12] = R4(A[13], 25); \
        B[13] = R4(A[19], 8); \
        B[14] = R4(A[20], 18); \
        B[15] = R4(A[4], 27); \
        B[16] = R4(A[5], 36); \
        B[17] = R4(A[11], 10); \
        B[18] = R4(A[17], 15); \
        B[19] = R4(A[23], 56); \
        B[20] = R4(A[2], 62); \
        B[21] = R4(A[8], 55); \
        B[22] = R4(A[14], 39); \
        B[23] = R4(A[15], 41); \
        B[24] = R4(A[21], 2); \
        \
        A[0] = _mm256_xor_si256(B[0], _mm256_andnot_si256(B[1], B[2])); \
        A[1] = _mm256_xor_si256(B[1], _mm256_andnot_si256(B[2], B[3])); \
        A[2] = _mm256_xor_si256(B[2], _mm256_andnot_si256(B[3], B[4])); \
        A[3] = _mm256_xor_si256(B[3], _mm256_andnot_si256(B[4], B[0])); \
        A[4] = _mm256_xor_si256(B[4], _mm256_andnot_si256(B[0], B[1])); \
        \
        A[5] = _mm256_xor_si256(B[5], _mm256_andnot_si256(B[6], B[7])); \
        A[6] = _mm256_xor_si256(B[6], _mm256_andnot_si256(B[7], B[8])); \
        A[7] = _mm256_xor_si256(B[7], _mm256_andnot_si256(B[8], B[9])); \
        A[8] = _mm256_xor_si256(B[8], _mm256_andnot_si256(B[9], B[5])); \
        A[9] = _mm256_xor_si256(B[9], _mm256_andnot_si256(B[5], B[6])); \
        \
        A[10] = _mm256_xor_si256(B[10], _mm256_andnot_si256(B[11], B[12])); \
        A[11] = _mm256_xor_si256(B[11], _mm256_andnot_si256(B[12], B[13])); \
        A[12] = _mm256_xor_si256(B[12], _mm256_andnot_si256(B[13], B[14])); \
        A[13] = _mm256_xor_si256(B[13], _mm256_andnot_si256(B[14], B[10])); \
        A[14] = _mm256_xor_si256(B[14], _mm256_andnot_si256(B[10], B[11])); \
        \
        A[15] = _mm256_xor_si256(B[15], _mm256_andnot_si256(B[16], B[17])); \
        A[16] = _mm256_xor_si256(B[16], _mm256_andnot_si256(B[17], B[18])); \
        A[17] = _mm256_xor_si256(B[17], _mm256_andnot_si256(B[18], B[19])); \
        A[18] = _mm256_xor_si256(B[18], _mm256_andnot_si256(B[19], B[15])); \
        A[19] = _mm256_xor_si256(B[19], _mm256_andnot_si256(B[15], B[16])); \
        \
        A[20] = _mm256_xor_si256(B[20], _mm256_andnot_si256(B[21], B[22])); \
        A[21] = _mm256_xor_si256(B[21], _mm256_andnot_si256(B[22], B[23])); \
        A[22] = _mm256_xor_si256(B[22], _mm256_andnot_si256(B[23], B[24])); \
        A[23] = _mm256_xor_si256(B[23], _mm256_andnot_si256(B[24], B[20])); \
        A[24] = _mm256_xor_si256(B[24], _mm256_andnot_si256(B[20], B[21])); \
        \
        A[0] = _mm256_xor_si256( \
            A[0], \
            _mm256_set1_epi64x(static_cast<int64_t>(RC[i])) \
        ); \
    } while (0)

static void keccak_f_4way(__m256i A[25]) {
    __m256i C[5];
    __m256i D[5];
    __m256i B[25];

    /*
     * DeepSeekHashV1 использует раунды 1..23.
     * Стандартный Keccak использовал бы также RC[0].
     */
    KF_ROUND4(1);
    KF_ROUND4(2);
    KF_ROUND4(3);
    KF_ROUND4(4);
    KF_ROUND4(5);
    KF_ROUND4(6);
    KF_ROUND4(7);
    KF_ROUND4(8);
    KF_ROUND4(9);
    KF_ROUND4(10);
    KF_ROUND4(11);
    KF_ROUND4(12);
    KF_ROUND4(13);
    KF_ROUND4(14);
    KF_ROUND4(15);
    KF_ROUND4(16);
    KF_ROUND4(17);
    KF_ROUND4(18);
    KF_ROUND4(19);
    KF_ROUND4(20);
    KF_ROUND4(21);
    KF_ROUND4(22);
    KF_ROUND4(23);
}

#undef KF_ROUND4
#undef R4

static void worker_avx2(
    const PowCtx& ctx,
    uint64_t from,
    uint64_t to,
    std::atomic<int64_t>& result
) {
    const __m256i target0 = _mm256_set1_epi64x(
        static_cast<int64_t>(ctx.target4[0])
    );

    const __m256i target1 = _mm256_set1_epi64x(
        static_cast<int64_t>(ctx.target4[1])
    );

    const __m256i target2 = _mm256_set1_epi64x(
        static_cast<int64_t>(ctx.target4[2])
    );

    const __m256i target3 = _mm256_set1_epi64x(
        static_cast<int64_t>(ctx.target4[3])
    );

    alignas(32) uint8_t dynamic[4][DYN_SIZE];
    char nonce_buffer[4][21];

    for (uint64_t nonce = from; nonce < to; nonce += 4) {
        if (
            (nonce & 0xFFFu) == 0 &&
            result.load(std::memory_order_relaxed) >= 0
        ) {
            return;
        }

        uint64_t lanes = to - nonce;

        if (lanes > 4) {
            lanes = 4;
        }

        for (uint64_t lane = 0; lane < lanes; ++lane) {
            std::memcpy(
                dynamic[lane],
                ctx.dyn_tpl,
                sizeof(ctx.dyn_tpl)
            );

            const int nonce_len = fast_u64_to_dec(
                nonce + lane,
                nonce_buffer[lane]
            );

            const size_t padding_position =
                ctx.dyn_offset +
                static_cast<size_t>(nonce_len);

            if (padding_position >= DYN_SIZE) {
                return;
            }

            std::memcpy(
                dynamic[lane] + ctx.dyn_offset,
                nonce_buffer[lane],
                static_cast<size_t>(nonce_len)
            );

            dynamic[lane][padding_position] = 0x06;
        }

        for (uint64_t lane = lanes; lane < 4; ++lane) {
            std::memcpy(
                dynamic[lane],
                dynamic[0],
                sizeof(dynamic[0])
            );
        }

        alignas(32) __m256i state[25];

        for (int word = 0; word < 25; ++word) {
            state[word] = _mm256_set1_epi64x(
                static_cast<int64_t>(
                    ctx.static_state[word]
                )
            );
        }

        for (
            size_t word = 0;
            word < ctx.dyn_word_count;
            ++word
        ) {
            uint64_t value0;
            uint64_t value1;
            uint64_t value2;
            uint64_t value3;

            std::memcpy(
                &value0,
                &dynamic[0][word * 8],
                sizeof(value0)
            );

            std::memcpy(
                &value1,
                &dynamic[1][word * 8],
                sizeof(value1)
            );

            std::memcpy(
                &value2,
                &dynamic[2][word * 8],
                sizeof(value2)
            );

            std::memcpy(
                &value3,
                &dynamic[3][word * 8],
                sizeof(value3)
            );

            const __m256i values = _mm256_set_epi64x(
                static_cast<int64_t>(value3),
                static_cast<int64_t>(value2),
                static_cast<int64_t>(value1),
                static_cast<int64_t>(value0)
            );

            state[
                ctx.dyn_word_start + word
            ] = _mm256_xor_si256(
                state[ctx.dyn_word_start + word],
                values
            );
        }

        keccak_f_4way(state);

        int mask = _mm256_movemask_pd(
            _mm256_castsi256_pd(
                _mm256_cmpeq_epi64(
                    state[0],
                    target0
                )
            )
        );

        if (!mask) {
            continue;
        }

        mask &= _mm256_movemask_pd(
            _mm256_castsi256_pd(
                _mm256_cmpeq_epi64(
                    state[1],
                    target1
                )
            )
        );

        if (!mask) {
            continue;
        }

        mask &= _mm256_movemask_pd(
            _mm256_castsi256_pd(
                _mm256_cmpeq_epi64(
                    state[2],
                    target2
                )
            )
        );

        if (!mask) {
            continue;
        }

        mask &= _mm256_movemask_pd(
            _mm256_castsi256_pd(
                _mm256_cmpeq_epi64(
                    state[3],
                    target3
                )
            )
        );

        if (!mask) {
            continue;
        }

        mask &= static_cast<int>(
            (1u << lanes) - 1u
        );

        for (
            int lane = 0;
            lane < static_cast<int>(lanes);
            ++lane
        ) {
            if (!(mask & (1 << lane))) {
                continue;
            }

            int64_t expected = -1;

            result.compare_exchange_strong(
                expected,
                static_cast<int64_t>(nonce) + lane,
                std::memory_order_relaxed
            );

            return;
        }
    }
}

static int64_t solve_internal(
    const std::string& base,
    const std::string& challenge_hex,
    int64_t difficulty
) {
    if (!valid_hex_digest(challenge_hex)) {
        return -3;
    }

    if (difficulty <= 0) {
        return -3;
    }

    /*
     * Solver рассчитан на один SHA3-256 rate block.
     */
    if (base.size() > RATE - MAX_NONCE_DEC - 1) {
        return -3;
    }

    PowCtx ctx = build_ctx(
        base,
        challenge_hex
    );

    std::atomic<int64_t> result{-1};

    unsigned thread_count =
        std::thread::hardware_concurrency();

    if (thread_count == 0) {
        thread_count = 1;
    }

    if (
        static_cast<int64_t>(thread_count) >
        difficulty
    ) {
        thread_count =
            static_cast<unsigned>(difficulty);
    }

    const uint64_t unsigned_difficulty =
        static_cast<uint64_t>(difficulty);

    const uint64_t chunk =
        (
            unsigned_difficulty +
            thread_count -
            1
        ) / thread_count;

    std::vector<std::thread> threads;
    threads.reserve(thread_count);

    for (
        unsigned thread = 0;
        thread < thread_count;
        ++thread
    ) {
        const uint64_t from =
            static_cast<uint64_t>(thread) *
            chunk;

        const uint64_t to = std::min(
            from + chunk,
            unsigned_difficulty
        );

        if (from >= to) {
            break;
        }

        threads.emplace_back(
            worker_avx2,
            std::cref(ctx),
            from,
            to,
            std::ref(result)
        );
    }

    for (auto& thread : threads) {
        thread.join();
    }

    return result.load(
        std::memory_order_relaxed
    );
}

extern "C" int64_t deepseek_pow_solve(
    const char* base,
    size_t base_len,
    const char* challenge_hex,
    size_t challenge_hex_len,
    int64_t difficulty
) {
    if (
        base == nullptr ||
        challenge_hex == nullptr
    ) {
        return -2;
    }

    try {
        const std::string base_string(
            base,
            base_len
        );

        const std::string challenge_string(
            challenge_hex,
            challenge_hex_len
        );

        return solve_internal(
            base_string,
            challenge_string,
            difficulty
        );
    } catch (const std::exception&) {
        return -4;
    } catch (...) {
        return -4;
    }
}