#include "BigInt.h"
#include "MemoryManager.h"
#include "Types.h"
#include <limits.h>
#include <stdio.h>
#include <string.h>

#define DIVIDE_UP(x, y) (1 + ((x) - 1) / (y))

typedef uint_fast32_t UWord;
#define UWORD_MAX UINT_FAST32_MAX
#define UWORD_BIT (sizeof(UWord) * CHAR_BIT)

/** Metadata **/

struct BigIntMeta {
    MemoryManager *mm;
    usize words;
    int resizeable;
};

static inline struct BigIntMeta*
B_metadata(const BigInt s) {
    return (struct BigIntMeta*)(s - sizeof(struct BigIntMeta));
}

/** Helpers **/

static inline void*
B_start(MemoryManager *mm, usize words) {
    if (NULL == mm)
        return NULL;

    struct BigIntMeta bim = {
        .mm = mm,
        .words = words,
        .resizeable = 1,
    };

    void* start = NULL;
    if (R_Ok != mm->alloc(*mm, &start, sizeof(bim) + bim.words * sizeof(UWord)) ||
        NULL == start)
    {
        return NULL;
    }

    memcpy(start, &bim, sizeof(bim));
    return start + sizeof(bim);
}

static inline Result
B_resize(BigInt *x, usize words) {
    struct BigIntMeta *bim = B_metadata(*x);

    void *start = bim;
    RETURN_NOTOK(
        bim->mm->realloc(*bim->mm, &start, sizeof(*bim) + words * sizeof(UWord)));

    bim = start;
    *x = start + sizeof(struct BigIntMeta);
    bim->words = words;

    return R_Ok;
}

/** Create **/

BigInt
B_new(MemoryManager *mm, int value) {
    UWord *start = B_start(mm, 1);

    if (NULL != start)
        memcpy(start, &value, sizeof(value));

    return start;
}

BigInt
B_newstr(MemoryManager *mm, const char* str) {
    if (NULL == str)
        return NULL;

    usize len = strlen(str);

    UWord *num = B_start(mm, DIVIDE_UP(len * CHAR_BIT / 2, UWORD_BIT));
    if (NULL == num)
        return NULL;

    struct BigIntMeta *numm = B_metadata(num);

    UWord *bcd = B_start(mm, numm->words);
    if (NULL == bcd) {
        mm->free(*mm, (void**)&num);
        return NULL;
    }

    struct BigIntMeta *bcdm = B_metadata(bcd);

    // char string to BCD
    for (const char *s = str; *s != '\0'; ++s) {
        B_lshift((BigInt*)&bcd, 4);
        bcd[0] |= *s - '0';
    }

    // Double-dable bcd to BigInt
    for (usize i = len * 4; i > 0; --i) {
        B_rshift((BigInt*)&num, 1);
        num[numm->words - 1] |= (UWord)(bcd[0] & 1) << (UWORD_BIT - 1);
        B_rshift((BigInt*)&bcd, 1);

        for (usize w = 0; w < bcdm->words; ++w) {
            for (UWord mask = 0b1111, eight = 8, three = 3; mask > 0; mask <<= 4, eight <<= 4, three <<= 4) {
                if ((bcd[w] & mask) >= eight)
                    bcd[w] -= three;
            }
        }
    }
    B_rshift((BigInt*)&num, UWORD_BIT * numm->words - len * 4);

    mm->free(*mm, (void**)&bcd);

    return num;
}

BigInt
B_newbytes(MemoryManager *mm, usize bytes) {
    if (0 == bytes)
        return B_start(mm, 1);
    return B_start(mm, DIVIDE_UP(bytes, sizeof(UWord)));
}

BigInt
B_dup(MemoryManager *mm, const BigInt bigint) {
    if (NULL == bigint)
        return NULL;

    struct BigIntMeta *bim = B_metadata(bigint);

    void *start = B_start(mm, bim->words);

    if (NULL != start)
        memcpy(start, bigint, bim->words * sizeof(UWord));

    return start;
}

/* Free */

Result
B_free(BigInt *bigint) {
    if (NULL == bigint || NULL == *bigint)
        return R_NullArgument;

    struct BigIntMeta *start = B_metadata(*bigint);
    *bigint = (BigInt*)start;
    return start->mm->free(*start->mm, (void*)bigint);
}

/* Modify */

Result
B_sum(BigInt *inta, const BigInt intb) {
    if (NULL == inta || NULL == intb)
        return R_NullArgument;

    if (NULL == *inta)
        return R_InvalidArgument;

    struct BigIntMeta *am = B_metadata(*inta), *bm = B_metadata(intb);

    if (am->resizeable && am->words < bm->words) {
        RETURN_NOTOK(B_resize(inta, bm->words));
        am = B_metadata(*inta);
    }

    UWord *a = *inta, *b = intb;

    int overflow = 0;
    usize words = am->words < bm->words ? am->words : bm->words;

    for (usize w = 0; w < words; ++w) {
        if (!(overflow && a[w] == UWORD_MAX)) {
            if (overflow)
                ++a[w];
            overflow = a[w] > UWORD_MAX - b[w];
        }

        a[w] += b[w];
    }

    if (overflow) {
        if (am->words > words) {
            for (usize w = words; w < am->words && overflow; ++w) {
                overflow = a[w] == UWORD_MAX;
                ++a[w];
            }
        }
        else if (am->resizeable) {
            RETURN_NOTOK(B_resize(inta, am->words << 1));
            a = *inta;
            ++a[words];
        }
    }

    return R_Ok;
}

Result
B_sub(BigInt *inta, const BigInt intb) {
    if (NULL == inta || NULL == intb)
        return R_NullArgument;

    if (NULL == *inta)
        return R_InvalidArgument;

    struct BigIntMeta *am = B_metadata(*inta), *bm = B_metadata(intb);

    if (am->resizeable && am->words < bm->words) {
        RETURN_NOTOK(B_resize(inta, bm->words));
        am = B_metadata(*inta);
    }

    UWord *a = *inta, *b = intb;

    int underflow = 0;
    usize words = am->words < bm->words ? am->words : bm->words;

    for (usize w = 0; w < words; ++w) {
        if (!(underflow && a[w] == 0)) {
            if (underflow)
                --a[w];
            underflow = a[w] < b[w];
        }

        a[w] -= b[w];
    }

    if (underflow) {
        if (am->words > words) {
            for (usize w = words; w < am->words && underflow; ++w) {
                underflow = a[w] == 0;
                --a[w];
            }
        }
        else if (am->resizeable) {
            RETURN_NOTOK(B_resize(inta, am->words << 1));
            a = *inta;
            --a[words];
        }
    }

    return R_Ok;
}



Result
B_compl(BigInt *intx) {
    if (NULL == intx)
        return R_NullArgument;

    struct BigIntMeta *xm = B_metadata(intx);
    UWord *x = (void*)intx;

    for (usize w = 0; w < xm->words; ++w)
        x[w] = ~x[w];

    return R_Ok;
}

Result
B_bitand(BigInt *intx, const BigInt inty) {
    if (NULL == intx || NULL == inty)
        return R_NullArgument;

    if (NULL == *intx)
        return R_InvalidArgument;

    struct BigIntMeta *xm = B_metadata(*intx), *ym = B_metadata(inty);

    if (xm->resizeable && xm->words < ym->words) {
        RETURN_NOTOK(B_resize(intx, ym->words));
        xm = B_metadata(intx);
    }

    UWord *x = (void*)intx, *y = (void*)inty;

    usize words = xm->words < ym->words ? xm->words : ym->words;

    for (usize w = 0; w < words; ++w)
        x[w] &= y[w];

    return R_Ok;
}

Result
B_bitor(BigInt *intx, const BigInt inty) {
    if (NULL == intx || NULL == inty)
        return R_NullArgument;

    if (NULL == *intx)
        return R_InvalidArgument;

    struct BigIntMeta *xm = B_metadata(*intx), *ym = B_metadata(inty);

    if (xm->resizeable && xm->words < ym->words) {
        RETURN_NOTOK(B_resize(intx, ym->words));
        xm = B_metadata(intx);
    }

    UWord *x = (void*)intx, *y = (void*)inty;

    usize words = xm->words < ym->words ? xm->words : ym->words;

    for (usize w = 0; w < words; ++w)
        x[w] |= y[w];

    return R_Ok;
}

Result
B_bitxor(BigInt *intx, const BigInt inty) {
    if (NULL == intx || NULL == inty)
        return R_NullArgument;

    if (NULL == *intx)
        return R_InvalidArgument;

    struct BigIntMeta *xm = B_metadata(*intx), *ym = B_metadata(inty);

    if (xm->resizeable && xm->words < ym->words) {
        RETURN_NOTOK(B_resize(intx, ym->words));
        xm = B_metadata(intx);
    }

    UWord *x = (void*)intx, *y = (void*)inty;

    usize words = xm->words < ym->words ? xm->words : ym->words;

    for (usize w = 0; w < words; ++w)
        x[w] ^= y[w];

    return R_Ok;
}

Result
B_rshift(BigInt *intx, usize shift) {
    if (NULL == intx)
        return R_NullArgument;

    if (NULL == *intx)
        return R_InvalidArgument;

    if (0 == shift)
        return R_Ok;

    struct BigIntMeta *xm = B_metadata(*intx);

    if (0 == xm->words)
        return R_Uninitialized;

    UWord *x = *intx;

    usize word_moves = shift / UWORD_BIT;
    if (word_moves) {
        usize w = 0;
        for (; w < xm->words - word_moves; ++w)
            x[w] = x[w + word_moves];
        for (; w < xm->words; ++w)
            x[w] = 0;
    }

    usize bit_moves = shift % UWORD_BIT;
    if (bit_moves == 1) {
        x[0] >>= 1;
        for (usize w = 1; w < xm->words; ++w) {
            x[w - 1] |= (x[w] & 1) << (UWORD_BIT - 1);
            x[w] >>= 1;
        }
    }
    else if (bit_moves > 1) {
        UWord mask = ((UWord)1 << bit_moves) - 1;

        x[0] >>= bit_moves;
        for (usize w = 1; w < xm->words; ++w) {
            x[w - 1] |= (x[w] & mask) << (UWORD_BIT - bit_moves);
            x[w] >>= bit_moves;
        }
    }

    return R_Ok;
}

Result
B_lshift(BigInt *intx, usize shift) {
    if (NULL == intx)
        return R_NullArgument;

    if (NULL == *intx)
        return R_InvalidArgument;

    if (0 == shift)
        return R_Ok;

    struct BigIntMeta *xm = B_metadata(*intx);

    if (0 == xm->words)
        return R_Uninitialized;

    if (xm->resizeable) {
        UWord width = B_bitwidth(*intx);
        if (width + shift > xm->words * UWORD_BIT) {
            RETURN_NOTOK(B_resize(intx, DIVIDE_UP(width + shift, UWORD_BIT)));
            xm = B_metadata(*intx);
        }
    }

    UWord *x = *intx;

    usize word_moves = shift / UWORD_BIT;
    if (word_moves) {
        usize w = xm->words - 1;
        for (; w >= word_moves; --w)
            x[w] = x[w - word_moves];
        for (; w < xm->words; --w)
            x[w] = 0;
    }

    usize bit_moves = shift % UWORD_BIT;
    if (bit_moves == 1) {
        UWord mask = (UWord)1 << (UWORD_BIT - 1);

        x[xm->words - 1] <<= 1;
        for (usize w = xm->words - 2; w < xm->words; --w) {
            x[w + 1] |= (x[w] & mask) >> (UWORD_BIT - 1);
            x[w] <<= 1;
        }
    }
    else if (bit_moves > 1) {
        UWord mask = (((UWord)1 << bit_moves) - 1) << (UWORD_BIT - bit_moves);

        x[xm->words - 1] <<= bit_moves;
        for (usize w = xm->words - 2; w < xm->words; --w) {
            x[w + 1] |= (x[w] & mask) >> (UWORD_BIT - bit_moves);
            x[w] <<= bit_moves;
        }
    }

    return R_Ok;
}

/* Get */

usize
B_bits(const BigInt x) {
    if (NULL == x)
        return 0;
    return B_metadata(x)->words * UWORD_BIT;
}

usize
B_bitwidth(const BigInt intx) {
    if (NULL == intx)
        return 0;

    struct BigIntMeta *bim = B_metadata(intx);
    UWord *x = intx;

    for (usize w = bim->words - 1; w < bim->words; --w) {
        if (x[w] == 0) continue;

        // TODO: use faster search than linear
        usize width = UWORD_BIT * (w + 1);
        for (UWord mask = (UWord)1 << (UWORD_BIT - 1); mask > 0; mask >>= 1, --width) {
            if (x[w] & mask)
                return width;
        }
    }

    return 0;
}

char*
B_tostr(MemoryManager *mm, const BigInt intx) {
    if (NULL == intx)
        return NULL;

    struct BigIntMeta *bim = B_metadata(intx);
    UWord *x = intx;

    usize n = B_bitwidth(intx);

    if (10 > n) {
        char *str = NULL;
        if (R_Ok != mm->alloc(*mm, (void**)&str, sizeof(char) * 2))
            return NULL;

        str[0] = '0' + n;
        str[1] = '\0';
        return str;
    }

    // The exact amount of necessary bcd bits is "ceil(log10(2^x - 1))",
    // however this is difficult to compute.
    // Wikipedia offers a convenient approximation "4 x ceil(n / 3)",
    // which grows faster than the true formula. Simplifying and calculating
    // for number of bytes, results in the following formula value.
    // Do note, a side effect is that we'll use more memory than needed and
    // show unnecessary leading 0s. Later in the code we'll discard these zeroes.
    usize bcd_bytes = DIVIDE_UP(n, 6);

    u8 *bcd;
    if (R_Ok != mm->alloc(*mm, (void**)&bcd, sizeof(u8) * bcd_bytes))
        return NULL;

    memset(bcd, 0, sizeof(u8) * bcd_bytes);

    UWord mask = (UWord)1 << (UWORD_BIT - 1);

    // Double-dable, binary to binary-coded decimal
    for (usize w = bim->words - 1; w < bim->words; --w) {
        for (UWord m = mask; m != 0; m >>= 1) {
            // Incrementing values >=5 by 3
            for (usize b = 0; b < bcd_bytes; ++b) {
                // Upper nibble
                if ((bcd[b] & 0b11110000) >= 0b01010000)
                    bcd[b] += 0b00110000;
                // Lower nibble
                if ((bcd[b] & 0b00001111) >= 5)
                    bcd[b] += 3;
            }

            // Left shift all bits in bcd
            u8 carry = 0, next_carry = 0;
            for (usize b = 0; b < bcd_bytes; ++b) {
                next_carry = !!(bcd[b] & 0b10000000);
                bcd[b] <<= 1;
                bcd[b] |= carry;
                carry = next_carry;
            }

            // Place leading bit in x into bcd
            bcd[0] |= !!(x[w] & m);
        }
    }

    // Remove leading zeroes
    while (bcd_bytes > 1 && bcd[bcd_bytes - 1] == 0)
        --bcd_bytes;

    int skip_nibble = (bcd[bcd_bytes - 1] & 0b11110000) == 0;

    char *str;
    if (R_Ok != mm->alloc(*mm, (void**)&str, sizeof(char) * (bcd_bytes * 2 - skip_nibble + 1))) {
        mm->free(*mm, (void**)&bcd);
        return NULL;
    }

    // BCD to char string
    char *s = str;

    if (!skip_nibble)
        *(s++) = ((bcd[bcd_bytes - 1] & 0b11110000) >> 4) + '0';
    *(s++)     =  (bcd[bcd_bytes - 1] & 0b00001111)       + '0';

    for (usize b = bcd_bytes - 2; b < bcd_bytes; --b) {
        *(s++) = ((bcd[b] & 0b11110000) >> 4) + '0';
        *(s++) =  (bcd[b] & 0b00001111)       + '0';
    }
    *s = '\0';

    mm->free(*mm, (void**)&bcd);
    return str;
}

int
B_eq(const BigInt intx, const BigInt inty) {
    if (NULL == intx || NULL == inty)
        return intx == inty;

    struct BigIntMeta *xm = B_metadata(intx), *ym = B_metadata(inty);
    UWord *x = intx, *y = inty;

    usize words;
    if (xm->words < ym->words) {
        for (usize w = xm->words; w < ym->words; ++w)
            if (y[w] != 0)
                return 0;

        words = xm->words;
    }
    else {
        for (usize w = ym->words; w < xm->words; ++w)
            if (x[w] != 0)
                return 0;

        words = ym->words;
    }

    for (usize w = words - 1; w < words; --w)
        if (x[w] != y[w])
            return 0;

    return 1;
}

int
B_not_eq(const BigInt x, const BigInt y) {
    return !B_eq(x, y);
}

int
B_lt(const BigInt intx, const BigInt inty) {
    if (NULL == intx || NULL == inty)
        return R_NullArgument;

    struct BigIntMeta *xm = B_metadata(intx), *ym = B_metadata(inty);
    UWord *x = (void*)intx, *y = (void*)inty;

    usize words;
    if (xm->words < ym->words) {
        for (usize w = xm->words; w < ym->words; ++w)
            if (y[w] != 0)
                return 1;

        words = xm->words;
    }
    else {
        for (usize w = ym->words; w < xm->words; ++w)
            if (x[w] != 0)
                return 0;

        words = ym->words;
    }

    for (usize w = words - 1; w < words; --w) {
        if (x[w] < y[w])
            return 1;
        else if (x[w] > y[w])
            return 0;
    }

    return 0;
}

int
B_le(const BigInt x, const BigInt y) {
    return B_lt(x, y) || B_eq(x, y);
}

int
B_gt(const BigInt x, const BigInt y) {
    return !B_lt(x, y) && !B_eq(x, y);
}

int
B_ge(const BigInt x, const BigInt y) {
    return !B_lt(x, y);
}