#include <bh/internal/bigint.h>
#include <bh/unit.h>
#include <string.h>

static int round_trip(void)
{
    bh_bigint_t *data[10];
    size_t i;

    for (i = 0; i < 10; i++)
    {
        data[i] = bh_bigint_new();
        bh_unit_assert(data[i] != NULL);
    }

    bh_unit_assert(bh_bigint_set_int(data[0], -12345) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(data[0]) == -12345);

    bh_unit_assert(bh_bigint_set_int8(data[1], -120) == BH_OK);
    bh_unit_assert(bh_bigint_get_int8(data[1]) == -120);

    bh_unit_assert(bh_bigint_set_int16(data[2], -12345) == BH_OK);
    bh_unit_assert(bh_bigint_get_int16(data[2]) == -12345);

    bh_unit_assert(bh_bigint_set_int32(data[3], -12345678l) == BH_OK);
    bh_unit_assert(bh_bigint_get_int32(data[3]) == -12345678l);

    bh_unit_assert(bh_bigint_set_int64(data[4], -123456789012345ll) == BH_OK);
    bh_unit_assert(bh_bigint_get_int64(data[4]) == -123456789012345ll);

    bh_unit_assert(bh_bigint_set_uint(data[5], 12345) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint(data[5]) == 12345);

    bh_unit_assert(bh_bigint_set_uint8(data[6], 230) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint8(data[6]) == 230);

    bh_unit_assert(bh_bigint_set_uint16(data[7], 40000) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint16(data[7]) == 40000);

    bh_unit_assert(bh_bigint_set_uint32(data[8], 12345678ul) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint32(data[8]) == 12345678ul);

    bh_unit_assert(bh_bigint_set_uint64(data[9], 123456789012345ull) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(data[9]) == 123456789012345ull);

    for (i = 0; i < 10; i++)
        bh_bigint_free(data[i]);

    return 0;
}

static int addition(void)
{
    bh_bigint_t *a, *b, *c;
    bh_uint64_t ta, tb;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);

    ta = 0xA3B1DBC158176B96ull;
    tb = 0x2748D1622F3E90A2ull;

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_set_uint64(b, tb) == BH_OK);
    bh_unit_assert(bh_bigint_add(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(c) == (ta + tb));

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);

    return 0;
}

static int subtraction(void)
{
    bh_bigint_t *a, *b, *c;
    bh_uint64_t ta, tb;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);

    ta = 0xA3B1DBC158176B96ull;
    tb = 0x2748D1622F3E90A2ull;

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_set_uint64(b, tb) == BH_OK);
    bh_unit_assert(bh_bigint_sub(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(c) == (ta - tb));

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);

    return 0;
}

static int addition_sign(void)
{
    bh_bigint_t *a, *b, *c;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);

    bh_unit_assert(bh_bigint_set_int(a, 5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 3) == BH_OK);
    bh_unit_assert(bh_bigint_add(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 8);
    bh_unit_assert(bh_bigint_add(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 8);

    bh_unit_assert(bh_bigint_set_int(a, -5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 3) == BH_OK);
    bh_unit_assert(bh_bigint_add(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -2);
    bh_unit_assert(bh_bigint_add(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -2);

    bh_unit_assert(bh_bigint_set_int(a, 5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, -3) == BH_OK);
    bh_unit_assert(bh_bigint_add(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 2);
    bh_unit_assert(bh_bigint_add(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 2);

    bh_unit_assert(bh_bigint_set_int(a, -5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, -3) == BH_OK);
    bh_unit_assert(bh_bigint_add(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -8);
    bh_unit_assert(bh_bigint_add(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -8);

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);

    return 0;
}

static int subtraction_sign(void)
{
    bh_bigint_t *a, *b, *c;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);

    bh_unit_assert(bh_bigint_set_int(a, 5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 3) == BH_OK);
    bh_unit_assert(bh_bigint_sub(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 2);
    bh_unit_assert(bh_bigint_sub(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -2);

    bh_unit_assert(bh_bigint_set_int(a, -5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 3) == BH_OK);
    bh_unit_assert(bh_bigint_sub(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -8);
    bh_unit_assert(bh_bigint_sub(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 8);

    bh_unit_assert(bh_bigint_set_int(a, 5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, -3) == BH_OK);
    bh_unit_assert(bh_bigint_sub(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 8);
    bh_unit_assert(bh_bigint_sub(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -8);

    bh_unit_assert(bh_bigint_set_int(a, -5) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, -3) == BH_OK);
    bh_unit_assert(bh_bigint_sub(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == -2);
    bh_unit_assert(bh_bigint_sub(c, b, a) == BH_OK);
    bh_unit_assert(bh_bigint_get_int(c) == 2);

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);

    return 0;
}

static int multiplication(void)
{
    int i;
    bh_bigint_t *a, *b, *c, *d;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();
    d = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);
    bh_unit_assert(d != NULL);

    bh_unit_assert(bh_bigint_set_int(a, 1234567) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 1234567) == BH_OK);

    bh_unit_assert(bh_bigint_mul(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_set(d, a) == BH_OK);
    bh_unit_assert(bh_bigint_reserve(d, d->size + a->size + 1) == BH_OK);

    for (i = 1; i < 1234567; i++)
        bh_unit_assert(bh_bigint_add(d, d, a) == BH_OK);

    bh_unit_assert(bh_bigint_equal(c, d) == 0);
    bh_unit_assert(memcmp(c->data, d->data, c->size * sizeof(*c->data)) == 0);

    bh_unit_assert(bh_bigint_set_int(c, 6) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 2) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(d, 6) == BH_OK);

    for (i = 0; i < 250; i++)
    {
        bh_unit_assert(bh_bigint_mul(c, c, b) == BH_OK);
        bh_unit_assert(bh_bigint_add(d, d, d) == BH_OK);
    }

    bh_unit_assert(bh_bigint_equal(c, d) == 0);
    bh_unit_assert(memcmp(c->data, d->data, c->size * sizeof(*c->data)) == 0);

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);
    bh_bigint_free(d);

    return 0;
}

static int logic(void)
{
    bh_bigint_t *a, *b, *c;
    bh_uint64_t ta, tb;

    ta = 0xA3B1DBC158176B96ull;
    tb = 0x2748D1622F3E90A2ull;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_set_uint64(b, tb) == BH_OK);
    
    bh_unit_assert(bh_bigint_or(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(c) == (ta | tb));

    bh_unit_assert(bh_bigint_and(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(c) == (ta & tb));

    bh_unit_assert(bh_bigint_xor(c, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(c) == (ta ^ tb));

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);

    return 0;
}

static int shifts(void)
{
    bh_bigint_t *a, *b;
    bh_uint64_t ta, tb;

    ta = 0x73B1DBC158176B96ull;
    tb = 0x73B1DBC158176B90ull;

    a = bh_bigint_new();
    b = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);

    bh_unit_assert(bh_bigint_set_uint64(a, 0xFFull) == BH_OK);
    bh_unit_assert(bh_bigint_lsh(b, a, 4) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (0xFFull << 4));

    bh_unit_assert(bh_bigint_set_uint64(a, 0xFFull) == BH_OK);
    bh_unit_assert(bh_bigint_lsh(b, a, 15) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (0xFFull << 15));

    bh_unit_assert(bh_bigint_set_uint64(a, 0xFFull) == BH_OK);
    bh_unit_assert(bh_bigint_lsh(b, a, 31) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (0xFFull << 31));

    bh_unit_assert(bh_bigint_set_uint64(a, 0xFFull) == BH_OK);
    bh_unit_assert(bh_bigint_lsh(b, a, 35) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (0xFFull << 35));

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_lsh(b, a, 1) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (ta << 1));

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_rsh(b, a, 31) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (ta >> 31))

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_rsh(b, a, 15) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (ta >> 15))

    bh_unit_assert(bh_bigint_set_uint64(a, ta) == BH_OK);
    bh_unit_assert(bh_bigint_rsh(b, a, 4) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == (ta >> 4));

    bh_unit_assert(bh_bigint_lsh(b, b, 4) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == tb);

    bh_unit_assert(bh_bigint_rsh(b, a, 60) == BH_OK);
    bh_unit_assert(bh_bigint_get_uint64(b) == 0x7);
    
    bh_bigint_free(a);
    bh_bigint_free(b);

    return 0;
}

static int division(void)
{
    bh_bigint_t *a, *b, *c, *d;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();
    d = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);
    bh_unit_assert(d != NULL);

    bh_unit_assert(bh_bigint_set_uint64(a, 0x41F915031348A134ull) == BH_OK);
    bh_unit_assert(bh_bigint_set_uint64(b, 0x434115F6ull) == BH_OK);
    bh_unit_assert(bh_bigint_mul(a, a, b) == BH_OK);
    bh_unit_assert(bh_bigint_set_uint64(b, 0x153F14ABull) == BH_OK);

    bh_unit_assert(bh_bigint_divmod(c, d, a, b) == BH_OK);

    bh_unit_assert(bh_bigint_get_uint64(c) == 0xD0D561E41F7EA081ull);
    bh_unit_assert(bh_bigint_get_uint64(d) == 0x0459E1CDull);

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);
    bh_bigint_free(d);

    return 0;
}

static int power(void)
{
    bh_bigint_t *a, *b, *c, *d, *e;

    a = bh_bigint_new();
    b = bh_bigint_new();
    c = bh_bigint_new();
    d = bh_bigint_new();
    e = bh_bigint_new();

    bh_unit_assert(a != NULL);
    bh_unit_assert(b != NULL);
    bh_unit_assert(c != NULL);
    bh_unit_assert(d != NULL);
    bh_unit_assert(e != NULL);

    bh_unit_assert(bh_bigint_set_int(a, 137) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(b, 2048) == BH_OK);
    bh_unit_assert(bh_bigint_set_int(c, 100) == BH_OK);

    bh_unit_assert(bh_bigint_pow(d, a, 2048) == BH_OK);
    bh_unit_assert(bh_bigint_mod(d, d, c) == BH_OK);
    bh_unit_assert(bh_bigint_powm(e, a, b, c) == BH_OK);
    bh_unit_assert(bh_bigint_equal(d, e) == 0);
    bh_unit_assert(bh_bigint_get_int(d) == 21);

    bh_bigint_free(a);
    bh_bigint_free(b);
    bh_bigint_free(c);
    bh_bigint_free(d);
    bh_bigint_free(e);

    return 0;
}

int main(int argc, char **argv)
{
    (void)argc;
    (void)argv;

    /* Add unit tests */
    bh_unit_add("round_trip", round_trip);
    bh_unit_add("addition", addition);
    bh_unit_add("subtraction", subtraction);
    bh_unit_add("addition_sign", addition_sign);
    bh_unit_add("subtraction_sign", subtraction_sign);
    bh_unit_add("multiplication", multiplication);
    bh_unit_add("logic", logic);
    bh_unit_add("shifts", shifts);
    bh_unit_add("division", division);
    bh_unit_add("power", power);

    return bh_unit_run();
}