/*

 * Copyright (c) 2022, stelar7 <dudedbz@gmail.com>

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#include <AK/Endian.h>

#include <AK/Types.h>

#include <LibCrypto/Curves/Curve25519.h>

namespace Crypto::Curves {

void Curve25519::set(u32* state, u32 value)

{

    state[0] = value;

    for (auto i = 1; i < WORDS; i++) {

        state[i] = 0;

    }

}

void Curve25519::modular_square(u32* state, u32 const* value)

{

    // Compute R = (A ^ 2) mod p

    modular_multiply(state, value, value);

}

void Curve25519::modular_subtract(u32* state, u32 const* first, u32 const* second)

{

    // R = (A - B) mod p

    i64 temp = -19;

    for (auto i = 0; i < WORDS; i++) {

        temp += first[i];

        temp -= second[i];

        state[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    // Compute R = A + (2^255 - 19) - B

    state[7] += 0x80000000;

    modular_reduce(state, state);

}

void Curve25519::modular_add(u32* state, u32 const* first, u32 const* second)

{

    // R = (A + B) mod p

    u64 temp = 0;

    for (auto i = 0; i < WORDS; i++) {

        temp += first[i];

        temp += second[i];

        state[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    modular_reduce(state, state);

}

void Curve25519::modular_multiply(u32* state, u32 const* first, u32 const* second)

{

    // Compute R = (A * B) mod p

    u64 temp = 0;

    u64 carry = 0;

    u32 output[WORDS * 2];

    // Comba's method

    for (auto i = 0; i < 16; i++) {

        if (i < WORDS) {

            for (auto j = 0; j <= i; j++) {

                temp += (u64)first[j] * second[i - j];

                carry += temp >> 32;

                temp &= 0xFFFFFFFF;

            }

        } else {

            for (auto j = i - 7; j < WORDS; j++) {

                temp += (u64)first[j] * second[i - j];

                carry += temp >> 32;

                temp &= 0xFFFFFFFF;

            }

        }

        output[i] = temp & 0xFFFFFFFF;

        temp = carry & 0xFFFFFFFF;

        carry >>= 32;

    }

    // Reduce bit 255 (2^255 = 19 mod p)

    temp = (output[7] >> 31) * 19;

    // Mask the most significant bit

    output[7] &= 0x7FFFFFFF;

    // Fast modular reduction 1st pass

    for (auto i = 0; i < WORDS; i++) {

        temp += output[i];

        temp += (u64)output[i + 8] * 38;

        output[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    // Reduce bit 256 (2^256 = 38 mod p)

    temp *= 38;

    // Reduce bit 255 (2^255 = 19 mod p)

    temp += (output[7] >> 31) * 19;

    // Mask the most significant bit

    output[7] &= 0x7FFFFFFF;

    // Fast modular reduction 2nd pass

    for (auto i = 0; i < WORDS; i++) {

        temp += output[i];

        output[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    modular_reduce(state, output);

}

void Curve25519::export_state(u32* state, u8* output)

{

    for (u32 i = 0; i < WORDS; i++) {

        state[i] = AK::convert_between_host_and_little_endian(state[i]);

    }

    memcpy(output, state, BYTES);

}

void Curve25519::import_state(u32* state, u8 const* data)

{

    memcpy(state, data, BYTES);

    for (u32 i = 0; i < WORDS; i++) {

        state[i] = AK::convert_between_host_and_little_endian(state[i]);

    }

}

void Curve25519::modular_subtract_single(u32* r, u32 const* a, u32 b)

{

    i64 temp = -19;

    temp -= b;

    // Compute R = A - 19 - B

    for (u32 i = 0; i < 8; i++) {

        temp += a[i];

        r[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    // Compute R = A + (2^255 - 19) - B

    r[7] += 0x80000000;

    modular_reduce(r, r);

}

void Curve25519::modular_add_single(u32* state, u32 const* first, u32 second)

{

    u64 temp = second;

    // Compute R = A + B

    for (u32 i = 0; i < 8; i++) {

        temp += first[i];

        state[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    modular_reduce(state, state);

}

u32 Curve25519::modular_square_root(u32* r, u32 const* a, u32 const* b)

{

    u32 c[8];

    u32 u[8];

    u32 v[8];

    // To compute the square root of (A / B), the first step is to compute the candidate root x = (A / B)^((p+3)/8)

    modular_square(v, b);

    modular_multiply(v, v, b);

    modular_square(v, v);

    modular_multiply(v, v, b);

    modular_multiply(c, a, v);

    modular_square(u, c);

    modular_multiply(u, u, c);

    modular_square(u, u);

    modular_multiply(v, u, c);

    to_power_of_2n(u, v, 3);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, c);

    to_power_of_2n(u, v, 7);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, c);

    to_power_of_2n(u, v, 15);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, c);

    to_power_of_2n(u, v, 31);

    modular_multiply(v, u, v);

    to_power_of_2n(u, v, 62);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, c);

    to_power_of_2n(u, v, 125);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_square(u, u);

    modular_multiply(u, u, c);

    // The first candidate root is U = A * B^3 * (A * B^7)^((p - 5) / 8)

    modular_multiply(u, u, a);

    modular_square(v, b);

    modular_multiply(v, v, b);

    modular_multiply(u, u, v);

    // The second candidate root is V = U * sqrt(-1)

    modular_multiply(v, u, SQRT_MINUS_1);

    modular_square(c, u);

    modular_multiply(c, c, b);

    // Check whether B * U^2 = A

    u32 first_comparison = compare(c, a);

    modular_square(c, v);

    modular_multiply(c, c, b);

    // Check whether B * V^2 = A

    u32 second_comparison = compare(c, a);

    // Select the first or the second candidate root

    select(r, u, v, first_comparison);

    // Return 0 if the square root exists

    return first_comparison & second_comparison;

}

u32 Curve25519::compare(u32 const* a, u32 const* b)

{

    u32 mask = 0;

    for (u32 i = 0; i < 8; i++) {

        mask |= a[i] ^ b[i];

    }

    // Return 0 if A = B, else 1

    return ((u32)(mask | (~mask + 1))) >> 31;

}

void Curve25519::modular_reduce(u32* state, u32 const* data)

{

    // R = A mod p

    u64 temp = 19;

    u32 other[WORDS];

    for (auto i = 0; i < WORDS; i++) {

        temp += data[i];

        other[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    // Compute B = A - (2^255 - 19)

    other[7] -= 0x80000000;

    u32 mask = (other[7] & 0x80000000) >> 31;

    select(state, other, data, mask);

}

void Curve25519::to_power_of_2n(u32* state, u32 const* value, u8 n)

{

    // Pre-compute (A ^ 2) mod p

    modular_square(state, value);

    // Compute R = (A ^ (2^n)) mod p

    for (u32 i = 1; i < n; i++) {

        modular_square(state, state);

    }

}

void Curve25519::select(u32* state, u32 const* a, u32 const* b, u32 condition)

{

    // If B < (2^255 - 19) then R = B, else R = A

    u32 mask = condition - 1;

    for (auto i = 0; i < WORDS; i++) {

        state[i] = (a[i] & mask) | (b[i] & ~mask);

    }

}

void Curve25519::copy(u32* state, u32 const* value)

{

    for (auto i = 0; i < WORDS; i++) {

        state[i] = value[i];

    }

}

void Curve25519::modular_multiply_inverse(u32* state, u32 const* value)

{

    // Compute R = A^-1 mod p

    u32 u[WORDS];

    u32 v[WORDS];

    // Fermat's little theorem

    modular_square(u, value);

    modular_multiply(u, u, value);

    modular_square(u, u);

    modular_multiply(v, u, value);

    to_power_of_2n(u, v, 3);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, value);

    to_power_of_2n(u, v, 7);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, value);

    to_power_of_2n(u, v, 15);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, value);

    to_power_of_2n(u, v, 31);

    modular_multiply(v, u, v);

    to_power_of_2n(u, v, 62);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_multiply(v, u, value);

    to_power_of_2n(u, v, 125);

    modular_multiply(u, u, v);

    modular_square(u, u);

    modular_square(u, u);

    modular_multiply(u, u, value);

    modular_square(u, u);

    modular_square(u, u);

    modular_multiply(u, u, value);

    modular_square(u, u);

    modular_multiply(state, u, value);

}

void Curve25519::modular_multiply_single(u32* state, u32 const* first, u32 second)

{

    // Compute R = (A * B) mod p

    u64 temp = 0;

    u32 output[WORDS];

    for (auto i = 0; i < WORDS; i++) {

        temp += (u64)first[i] * second;

        output[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    // Reduce bit 256 (2^256 = 38 mod p)

    temp *= 38;

    // Reduce bit 255 (2^255 = 19 mod p)

    temp += (output[7] >> 31) * 19;

    // Mask the most significant bit

    output[7] &= 0x7FFFFFFF;

    // Fast modular reduction

    for (auto i = 0; i < WORDS; i++) {

        temp += output[i];

        output[i] = temp & 0xFFFFFFFF;

        temp >>= 32;

    }

    modular_reduce(state, output);

}

}