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

typedef uint_fast32_t UWord;
#define UWORD_MAX UINT_FAST32_MAX

/** Metadata **/

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

static inline struct BigIntMeta*
B_metadata(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,
    };

    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) {
    struct BigIntMeta *bim = B_metadata(*x);

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

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

    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) {
    return NULL;
}

BigInt
B_newbytes(MemoryManager *mm, usize bytes) {
    return B_start(mm, 1 + (bytes - 1) / sizeof(UWord));
}

BigInt
B_dup(MemoryManager *mm, 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;
}

BigInt
B_newsum(MemoryManager *mm, BigInt a, BigInt b) {
    BigInt total = B_dup(mm, a);

    if (NULL == total ||
        R_Ok != B_sum(&total, b))
    {
        return NULL;
    }

    return total;
}

BigInt
B_newsub(MemoryManager *mm, BigInt a, BigInt b) {
    BigInt total = B_dup(mm, a);

    if (NULL == total ||
        R_Ok != B_sub(&total, b))
    {
        return NULL;
    }

    return total;
}

/* 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, BigInt intb) {
    if (NULL == inta || NULL == intb || NULL == *inta)
        return R_NullArgument;

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

    int overflow = 0;
    usize w;
    for (w = 0; w < bm->words; ++w) {
        if (w >= am->words)
            B_resize(inta);

        if (!(overflow && a[w] > UWORD_MAX - 1)) {
            if (overflow)
                ++a[w];
            overflow = a[w] > UWORD_MAX - b[w];
        }

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

    if (overflow) {
        if (w >= am->words) {
            RETURN_NOTOK(B_resize(inta));
            a = *inta;
        }
        ++a[w];
    }

    return R_Ok;
}

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

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

    int underflow = 0;
    usize w;
    for (w = 0; w < bm->words; ++w) {
        if (w >= am->words)
            B_resize(inta);

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

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

    if (underflow) {
        if (w >= am->words) {
            RETURN_NOTOK(B_resize(inta));
            a = *inta;
        }
        --a[w];
    }

    return R_Ok;
}

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

    if (shift == 0)
        return R_Ok;

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

    usize word_moves = shift / (sizeof(UWord) * 8);
    if (word_moves) {
        usize w = 0;
        for (; w < bim->words - word_moves; ++w)
            x[w] = x[w + word_moves];
        for (; w < bim->words; ++w)
            x[w] = 0;
    }

    usize bit_moves = shift % (sizeof(UWord) * 8);
    if (bit_moves == 1) {
        x[0] >>= 1;
        for (usize w = 1; w < bim->words; ++w) {
            x[w - 1] |= (x[w] & 1) << (sizeof(UWord) * 8 - 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 < bim->words; ++w) {
            x[w - 1] |= (x[w] & mask) << (sizeof(UWord) * 8 - bit_moves);
            x[w] >>= bit_moves;
        }
    }

    return R_Ok;
}

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

    if (shift == 0)
        return R_Ok;

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

    usize word_moves = shift / (sizeof(UWord) * 8);
    if (word_moves) {
        usize w = bim->words - 1;
        for (; w >= word_moves; --w)
            x[w] = x[w - word_moves];
        for (; w < bim->words; --w)
            x[w] = 0;
    }

    usize bit_moves = shift % (sizeof(UWord) * 8);
    if (bit_moves == 1) {
        UWord mask = (UWord)1 << (sizeof(UWord) * 8 - 1);

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

        x[bim->words - 1] <<= bit_moves;
        for (usize w = bim->words - 2; w < bim->words; --w) {
            x[w + 1] |= (x[w] & mask) >> (sizeof(UWord) * 8 - bit_moves);
            x[w] <<= bit_moves;
        }
    }

    return R_Ok;
}

/* Get */

usize
B_bytes(BigInt x) {
    if (NULL == x)
        return 0;
    return B_metadata(x)->words * sizeof(UWord);
}

usize
B_bitwidth(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 = sizeof(UWord) * 8 * (w + 1);
        for (UWord mask = (UWord)1 << (sizeof(UWord) * 8 - 1); mask > 0; mask >>= 1, --width) {
            if (x[w] & mask)
                return width;
        }
    }

    return 0;
}

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

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

    usize n = B_bitwidth(intx);

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

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

    // TODO: these calculations are based on Wikipedia's statement that
    // "4 x ceil(n / 3)" is enough bits of memory.
    // This is true, but it grows faster than needed, so from 31 bits onward,
    // we use too much memory (adding unnecessary leading 0s).
    // The exact value is "ceil(log10(2^x - 1))"
    usize bcd_bytes = 1 + (n - 1) / 6;
    int odd_nibbles = (1 + (n - 1) / 3) % 2;

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

    u8 *bcd;
    if (R_Ok != mm->alloc(*mm, (void**)&bcd, sizeof(u8) * bcd_bytes)) {
        mm->free(*mm, (void**)&str);
        return NULL;
    }
    memset(bcd, 0, sizeof(u8) * bcd_bytes);

    UWord mask = (UWord)1 << (sizeof(UWord) * 8 - 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);
        }
    }

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

    if (!odd_nibbles)
        *(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;
}