#include "NativeXChaCha20.hpp"

#include <array>
#include <cstring>

namespace {

constexpr std::size_t kXChaChaKeySize = 32;
constexpr std::size_t kXChaChaNonceSize = 24;
constexpr std::size_t kChaChaNonceSize = 12;
constexpr std::size_t kChaChaBlockSize = 64;
constexpr std::size_t kPoly1305TagSize = 16;

inline std::uint32_t load_le32(const std::uint8_t *src) {
    return static_cast<std::uint32_t>(src[0])
        | (static_cast<std::uint32_t>(src[1]) << 8U)
        | (static_cast<std::uint32_t>(src[2]) << 16U)
        | (static_cast<std::uint32_t>(src[3]) << 24U);
}

inline std::uint64_t load_le64(const std::uint8_t *src) {
    return static_cast<std::uint64_t>(src[0])
        | (static_cast<std::uint64_t>(src[1]) << 8U)
        | (static_cast<std::uint64_t>(src[2]) << 16U)
        | (static_cast<std::uint64_t>(src[3]) << 24U)
        | (static_cast<std::uint64_t>(src[4]) << 32U)
        | (static_cast<std::uint64_t>(src[5]) << 40U)
        | (static_cast<std::uint64_t>(src[6]) << 48U)
        | (static_cast<std::uint64_t>(src[7]) << 56U);
}

inline void store_le32(std::uint32_t value, std::uint8_t *dst) {
    dst[0] = static_cast<std::uint8_t>(value & 0xFFU);
    dst[1] = static_cast<std::uint8_t>((value >> 8U) & 0xFFU);
    dst[2] = static_cast<std::uint8_t>((value >> 16U) & 0xFFU);
    dst[3] = static_cast<std::uint8_t>((value >> 24U) & 0xFFU);
}

inline void store_le64(std::uint64_t value, std::uint8_t *dst) {
    for (int i = 0; i < 8; i++) {
        dst[i] = static_cast<std::uint8_t>((value >> (i * 8U)) & 0xFFU);
    }
}

inline std::uint32_t rotl32(std::uint32_t value, int shift) {
    return (value << shift) | (value >> (32 - shift));
}

inline void quarter_round(std::uint32_t *state, int a, int b, int c, int d) {
    state[a] += state[b];
    state[d] ^= state[a];
    state[d] = (state[d] << 16U) | (state[d] >> 16U);

    state[c] += state[d];
    state[b] ^= state[c];
    state[b] = (state[b] << 12U) | (state[b] >> 20U);

    state[a] += state[b];
    state[d] ^= state[a];
    state[d] = (state[d] << 8U) | (state[d] >> 24U);

    state[c] += state[d];
    state[b] ^= state[c];
    state[b] = (state[b] << 7U) | (state[b] >> 25U);
}

void chacha20_block(
    const std::uint8_t key[kXChaChaKeySize],
    const std::uint8_t nonce[kChaChaNonceSize],
    std::uint32_t counter,
    std::uint8_t out[kChaChaBlockSize]
) {
    std::uint32_t state[16] = {};
    state[0] = 0x61707865U;
    state[1] = 0x3320646eU;
    state[2] = 0x79622d32U;
    state[3] = 0x6b206574U;

    for (int i = 0; i < 8; i++) {
        state[4 + i] = load_le32(key + (i * 4));
    }
    state[12] = counter;
    for (int i = 0; i < 3; i++) {
        state[13 + i] = load_le32(nonce + (i * 4));
    }

    std::uint32_t working[16];
    std::memcpy(working, state, sizeof(state));

    for (int i = 0; i < 10; i++) {
        quarter_round(working, 0, 4, 8, 12);
        quarter_round(working, 1, 5, 9, 13);
        quarter_round(working, 2, 6, 10, 14);
        quarter_round(working, 3, 7, 11, 15);
        quarter_round(working, 0, 5, 10, 15);
        quarter_round(working, 1, 6, 11, 12);
        quarter_round(working, 2, 7, 8, 13);
        quarter_round(working, 3, 4, 9, 14);
    }

    for (int i = 0; i < 16; i++) {
        working[i] += state[i];
        store_le32(working[i], out + (i * 4));
    }
}

void hchacha20(
    const std::uint8_t key[kXChaChaKeySize],
    const std::uint8_t nonce[16],
    std::uint8_t out[kXChaChaKeySize]
) {
    std::uint32_t state[16] = {};
    state[0] = 0x61707865U;
    state[1] = 0x3320646eU;
    state[2] = 0x79622d32U;
    state[3] = 0x6b206574U;

    for (int i = 0; i < 8; i++) {
        state[4 + i] = load_le32(key + (i * 4));
    }
    for (int i = 0; i < 4; i++) {
        state[12 + i] = load_le32(nonce + (i * 4));
    }

    for (int i = 0; i < 10; i++) {
        quarter_round(state, 0, 4, 8, 12);
        quarter_round(state, 1, 5, 9, 13);
        quarter_round(state, 2, 6, 10, 14);
        quarter_round(state, 3, 7, 11, 15);
        quarter_round(state, 0, 5, 10, 15);
        quarter_round(state, 1, 6, 11, 12);
        quarter_round(state, 2, 7, 8, 13);
        quarter_round(state, 3, 4, 9, 14);
    }

    for (int i = 0; i < 4; i++) {
        store_le32(state[i], out + (i * 4));
    }
    for (int i = 0; i < 4; i++) {
        store_le32(state[12 + i], out + (16 + i * 4));
    }
}

void chacha20_xor(
    const std::uint8_t *input,
    std::size_t input_length,
    const std::uint8_t key[kXChaChaKeySize],
    const std::uint8_t nonce[kChaChaNonceSize],
    std::uint32_t initial_counter,
    std::vector<std::uint8_t> &output
) {
    output.resize(input_length);
    std::uint32_t counter = initial_counter;
    std::uint8_t block[kChaChaBlockSize];

    for (std::size_t offset = 0; offset < input_length; offset += kChaChaBlockSize) {
        chacha20_block(key, nonce, counter, block);
        const std::size_t block_size = std::min(kChaChaBlockSize, input_length - offset);
        for (std::size_t i = 0; i < block_size; i++) {
            output[offset + i] = static_cast<std::uint8_t>(input[offset + i] ^ block[i]);
        }
        counter += 1;
    }
}

void hsalsa20_core(
    const std::uint8_t key[kXChaChaKeySize],
    const std::uint8_t nonce16[16],
    std::uint8_t out[kXChaChaKeySize]
) {
    std::uint32_t x[16] = {};
    x[0] = 0x61707865U;
    x[1] = load_le32(key + 0);
    x[2] = load_le32(key + 4);
    x[3] = load_le32(key + 8);
    x[4] = load_le32(key + 12);
    x[5] = 0x3320646eU;
    x[6] = load_le32(nonce16 + 0);
    x[7] = load_le32(nonce16 + 4);
    x[8] = load_le32(nonce16 + 8);
    x[9] = load_le32(nonce16 + 12);
    x[10] = 0x79622d32U;
    x[11] = load_le32(key + 16);
    x[12] = load_le32(key + 20);
    x[13] = load_le32(key + 24);
    x[14] = load_le32(key + 28);
    x[15] = 0x6b206574U;

    for (int i = 0; i < 20; i += 2) {
        // Salsa20 column round
        x[4] ^= rotl32(x[0] + x[12], 7);
        x[8] ^= rotl32(x[4] + x[0], 9);
        x[12] ^= rotl32(x[8] + x[4], 13);
        x[0] ^= rotl32(x[12] + x[8], 18);
        x[9] ^= rotl32(x[5] + x[1], 7);
        x[13] ^= rotl32(x[9] + x[5], 9);
        x[1] ^= rotl32(x[13] + x[9], 13);
        x[5] ^= rotl32(x[1] + x[13], 18);
        x[14] ^= rotl32(x[10] + x[6], 7);
        x[2] ^= rotl32(x[14] + x[10], 9);
        x[6] ^= rotl32(x[2] + x[14], 13);
        x[10] ^= rotl32(x[6] + x[2], 18);
        x[3] ^= rotl32(x[15] + x[11], 7);
        x[7] ^= rotl32(x[3] + x[15], 9);
        x[11] ^= rotl32(x[7] + x[3], 13);
        x[15] ^= rotl32(x[11] + x[7], 18);

        // Salsa20 row round
        x[1] ^= rotl32(x[0] + x[3], 7);
        x[2] ^= rotl32(x[1] + x[0], 9);
        x[3] ^= rotl32(x[2] + x[1], 13);
        x[0] ^= rotl32(x[3] + x[2], 18);
        x[6] ^= rotl32(x[5] + x[4], 7);
        x[7] ^= rotl32(x[6] + x[5], 9);
        x[4] ^= rotl32(x[7] + x[6], 13);
        x[5] ^= rotl32(x[4] + x[7], 18);
        x[11] ^= rotl32(x[10] + x[9], 7);
        x[8] ^= rotl32(x[11] + x[10], 9);
        x[9] ^= rotl32(x[8] + x[11], 13);
        x[10] ^= rotl32(x[9] + x[8], 18);
        x[12] ^= rotl32(x[15] + x[14], 7);
        x[13] ^= rotl32(x[12] + x[15], 9);
        x[14] ^= rotl32(x[13] + x[12], 13);
        x[15] ^= rotl32(x[14] + x[13], 18);
    }

    // NaCl box.before() output order: 0,5,10,15,6,7,8,9
    store_le32(x[0], out + 0);
    store_le32(x[5], out + 4);
    store_le32(x[10], out + 8);
    store_le32(x[15], out + 12);
    store_le32(x[6], out + 16);
    store_le32(x[7], out + 20);
    store_le32(x[8], out + 24);
    store_le32(x[9], out + 28);
}

void poly1305_to_limbs26(const std::uint8_t block16[16], std::uint64_t limbs[5]) {
    const std::uint64_t lo = load_le64(block16);
    const std::uint64_t hi = load_le64(block16 + 8);
    limbs[0] = lo & 0x3FFFFFFULL;
    limbs[1] = (lo >> 26U) & 0x3FFFFFFULL;
    limbs[2] = ((lo >> 52U) | (hi << 12U)) & 0x3FFFFFFULL;
    limbs[3] = (hi >> 14U) & 0x3FFFFFFULL;
    limbs[4] = (hi >> 40U) & 0x3FFFFFFULL;
}

void poly1305_multiply_reduce(std::uint64_t h[5], const std::uint64_t r[5]) {
    const std::uint64_t r0 = r[0];
    const std::uint64_t r1 = r[1];
    const std::uint64_t r2 = r[2];
    const std::uint64_t r3 = r[3];
    const std::uint64_t r4 = r[4];
    const std::uint64_t s1 = r1 * 5ULL;
    const std::uint64_t s2 = r2 * 5ULL;
    const std::uint64_t s3 = r3 * 5ULL;
    const std::uint64_t s4 = r4 * 5ULL;

    const std::uint64_t h0_in = h[0];
    const std::uint64_t h1_in = h[1];
    const std::uint64_t h2_in = h[2];
    const std::uint64_t h3_in = h[3];
    const std::uint64_t h4_in = h[4];

    std::uint64_t h0 = h0_in * r0 + h1_in * s4 + h2_in * s3 + h3_in * s2 + h4_in * s1;
    std::uint64_t h1 = h0_in * r1 + h1_in * r0 + h2_in * s4 + h3_in * s3 + h4_in * s2;
    std::uint64_t h2 = h0_in * r2 + h1_in * r1 + h2_in * r0 + h3_in * s4 + h4_in * s3;
    std::uint64_t h3 = h0_in * r3 + h1_in * r2 + h2_in * r1 + h3_in * r0 + h4_in * s4;
    std::uint64_t h4 = h0_in * r4 + h1_in * r3 + h2_in * r2 + h3_in * r1 + h4_in * r0;

    std::uint64_t c = h0 >> 26U;
    h1 += c;
    h0 &= 0x3FFFFFFULL;
    c = h1 >> 26U;
    h2 += c;
    h1 &= 0x3FFFFFFULL;
    c = h2 >> 26U;
    h3 += c;
    h2 &= 0x3FFFFFFULL;
    c = h3 >> 26U;
    h4 += c;
    h3 &= 0x3FFFFFFULL;
    c = h4 >> 26U;
    h0 += c * 5ULL;
    h4 &= 0x3FFFFFFULL;
    c = h0 >> 26U;
    h1 += c;
    h0 &= 0x3FFFFFFULL;

    h[0] = h0;
    h[1] = h1;
    h[2] = h2;
    h[3] = h3;
    h[4] = h4;
}

void poly1305_mac(
    const std::uint8_t *data,
    std::size_t data_length,
    const std::uint8_t key32[32],
    std::uint8_t out_tag16[16]
) {
    std::uint8_t r_bytes[16];
    std::memcpy(r_bytes, key32, 16);
    r_bytes[3] &= 15U;
    r_bytes[7] &= 15U;
    r_bytes[11] &= 15U;
    r_bytes[15] &= 15U;
    r_bytes[4] &= 252U;
    r_bytes[8] &= 252U;
    r_bytes[12] &= 252U;

    std::uint8_t s_bytes[16];
    std::memcpy(s_bytes, key32 + 16, 16);

    std::uint64_t r[5];
    poly1305_to_limbs26(r_bytes, r);
    std::uint64_t h[5] = {0, 0, 0, 0, 0};

    std::vector<std::uint8_t> mac_input;
    const std::size_t pad = (16 - (data_length % 16)) % 16;
    mac_input.reserve(data_length + pad + 16);
    mac_input.insert(mac_input.end(), data, data + data_length);
    mac_input.insert(mac_input.end(), pad, 0);
    mac_input.insert(mac_input.end(), 8, 0);
    std::uint8_t ct_len_le[8];
    store_le64(static_cast<std::uint64_t>(data_length), ct_len_le);
    mac_input.insert(mac_input.end(), ct_len_le, ct_len_le + 8);

    for (std::size_t offset = 0; offset < mac_input.size(); offset += 16) {
        std::uint8_t block[16];
        std::memcpy(block, mac_input.data() + offset, 16);

        std::uint64_t n[5];
        poly1305_to_limbs26(block, n);
        h[0] += n[0];
        h[1] += n[1];
        h[2] += n[2];
        h[3] += n[3];
        h[4] += n[4] + (1ULL << 24U);
        poly1305_multiply_reduce(h, r);
    }

    std::uint64_t h0 = h[0];
    std::uint64_t h1 = h[1];
    std::uint64_t h2 = h[2];
    std::uint64_t h3 = h[3];
    std::uint64_t h4 = h[4];

    std::uint64_t c = h0 >> 26U;
    h1 += c;
    h0 &= 0x3FFFFFFULL;
    c = h1 >> 26U;
    h2 += c;
    h1 &= 0x3FFFFFFULL;
    c = h2 >> 26U;
    h3 += c;
    h2 &= 0x3FFFFFFULL;
    c = h3 >> 26U;
    h4 += c;
    h3 &= 0x3FFFFFFULL;
    c = h4 >> 26U;
    h0 += c * 5ULL;
    h4 &= 0x3FFFFFFULL;
    c = h0 >> 26U;
    h1 += c;
    h0 &= 0x3FFFFFFULL;

    std::uint64_t g0 = h0 + 5ULL;
    c = g0 >> 26U;
    g0 &= 0x3FFFFFFULL;
    std::uint64_t g1 = h1 + c;
    c = g1 >> 26U;
    g1 &= 0x3FFFFFFULL;
    std::uint64_t g2 = h2 + c;
    c = g2 >> 26U;
    g2 &= 0x3FFFFFFULL;
    std::uint64_t g3 = h3 + c;
    c = g3 >> 26U;
    g3 &= 0x3FFFFFFULL;
    const std::uint64_t g4 = h4 + c - (1ULL << 26U);

    const std::uint64_t mask = (g4 >> 63U) - 1ULL;
    const std::uint64_t nmask = ~mask;
    h0 = (h0 & nmask) | (g0 & mask);
    h1 = (h1 & nmask) | (g1 & mask);
    h2 = (h2 & nmask) | (g2 & mask);
    h3 = (h3 & nmask) | (g3 & mask);
    h4 = (h4 & nmask) | (g4 & mask);

    const std::uint64_t lo = h0 | (h1 << 26U) | (h2 << 52U);
    const std::uint64_t hi = (h2 >> 12U) | (h3 << 14U) | (h4 << 40U);

    const std::uint64_t s_lo = load_le64(s_bytes);
    const std::uint64_t s_hi = load_le64(s_bytes + 8);
    const std::uint64_t result_lo = lo + s_lo;
    const std::uint64_t carry = (result_lo < lo) ? 1ULL : 0ULL;
    const std::uint64_t result_hi = hi + s_hi + carry;

    store_le64(result_lo, out_tag16);
    store_le64(result_hi, out_tag16 + 8);
}

bool constant_time_equal(const std::uint8_t *a, const std::uint8_t *b, std::size_t length) {
    std::uint8_t diff = 0;
    for (std::size_t i = 0; i < length; i++) {
        diff |= static_cast<std::uint8_t>(a[i] ^ b[i]);
    }
    return diff == 0;
}

void derive_subkey_and_nonce(
    const std::uint8_t key32[kXChaChaKeySize],
    const std::uint8_t nonce24[kXChaChaNonceSize],
    std::uint8_t out_subkey32[kXChaChaKeySize],
    std::uint8_t out_chacha_nonce12[kChaChaNonceSize]
) {
    hchacha20(key32, nonce24, out_subkey32);
    std::memset(out_chacha_nonce12, 0, kChaChaNonceSize);
    std::memcpy(out_chacha_nonce12 + 4, nonce24 + 16, 8);
}

} // namespace

namespace rosetta::nativecrypto {

bool xchacha20poly1305_encrypt(
    const std::uint8_t *plaintext,
    std::size_t plaintext_length,
    const std::uint8_t *key32,
    const std::uint8_t *nonce24,
    std::vector<std::uint8_t> &ciphertext_with_tag
) {
    if (key32 == nullptr || nonce24 == nullptr || (plaintext == nullptr && plaintext_length > 0)) {
        return false;
    }

    std::uint8_t subkey[kXChaChaKeySize];
    std::uint8_t chacha_nonce[kChaChaNonceSize];
    derive_subkey_and_nonce(key32, nonce24, subkey, chacha_nonce);

    std::vector<std::uint8_t> ciphertext;
    chacha20_xor(plaintext, plaintext_length, subkey, chacha_nonce, 1, ciphertext);

    std::uint8_t block[kChaChaBlockSize];
    chacha20_block(subkey, chacha_nonce, 0, block);
    std::uint8_t tag[kPoly1305TagSize];
    poly1305_mac(ciphertext.data(), ciphertext.size(), block, tag);

    ciphertext_with_tag = ciphertext;
    ciphertext_with_tag.insert(ciphertext_with_tag.end(), tag, tag + kPoly1305TagSize);
    return true;
}

bool xchacha20poly1305_decrypt(
    const std::uint8_t *ciphertext_with_tag,
    std::size_t ciphertext_with_tag_length,
    const std::uint8_t *key32,
    const std::uint8_t *nonce24,
    std::vector<std::uint8_t> &plaintext
) {
    if (key32 == nullptr || nonce24 == nullptr || ciphertext_with_tag == nullptr) {
        return false;
    }
    if (ciphertext_with_tag_length < kPoly1305TagSize) {
        return false;
    }

    const std::size_t ciphertext_length = ciphertext_with_tag_length - kPoly1305TagSize;
    const std::uint8_t *ciphertext = ciphertext_with_tag;
    const std::uint8_t *tag = ciphertext_with_tag + ciphertext_length;

    std::uint8_t subkey[kXChaChaKeySize];
    std::uint8_t chacha_nonce[kChaChaNonceSize];
    derive_subkey_and_nonce(key32, nonce24, subkey, chacha_nonce);

    std::uint8_t block[kChaChaBlockSize];
    chacha20_block(subkey, chacha_nonce, 0, block);
    std::uint8_t computed_tag[kPoly1305TagSize];
    poly1305_mac(ciphertext, ciphertext_length, block, computed_tag);
    if (!constant_time_equal(tag, computed_tag, kPoly1305TagSize)) {
        return false;
    }

    chacha20_xor(ciphertext, ciphertext_length, subkey, chacha_nonce, 1, plaintext);
    return true;
}

bool xchacha20_xor(
    const std::uint8_t *input,
    std::size_t input_length,
    const std::uint8_t *key32,
    const std::uint8_t *nonce24,
    std::vector<std::uint8_t> &output
) {
    if (key32 == nullptr || nonce24 == nullptr || (input == nullptr && input_length > 0)) {
        return false;
    }

    std::uint8_t subkey[kXChaChaKeySize];
    std::uint8_t chacha_nonce[kChaChaNonceSize];
    derive_subkey_and_nonce(key32, nonce24, subkey, chacha_nonce);

    // Stream-cipher mode parity with desktop/libsodium crypto_stream_xchacha20_xor.
    chacha20_xor(input, input_length, subkey, chacha_nonce, 0, output);

    std::memset(subkey, 0, sizeof(subkey));
    std::memset(chacha_nonce, 0, sizeof(chacha_nonce));
    return true;
}

bool hsalsa20_derive(
    const std::uint8_t *raw_dh32,
    std::vector<std::uint8_t> &shared_key32
) {
    if (raw_dh32 == nullptr) {
        return false;
    }

    const std::uint8_t zero_nonce[16] = {0};
    std::uint8_t out[kXChaChaKeySize];
    hsalsa20_core(raw_dh32, zero_nonce, out);
    shared_key32.assign(out, out + kXChaChaKeySize);
    std::memset(out, 0, sizeof(out));
    return true;
}

} // namespace rosetta::nativecrypto
