#include "alloc.h"
#include "klibc.h"
#include "serial.h"
#include "vga.h"

int memcmp(const void *aptr, const void *bptr, size_t size)
{
    const unsigned char *a = (const unsigned char *)aptr;
    const unsigned char *b = (const unsigned char *)bptr;
    for (size_t i = 0; i < size; i++) {
        if (a[i] < b[i])
            return -1;
        else if (b[i] < a[i])
            return 1;
    }
    return 0;
}

// Inspirated by https://interrupt.memfault.com/blog/memcpy-newlib-nano
/* Nonzero if either X or Y is not aligned on a "long" boundary.  */
#define UNALIGNED(X, Y) \
  (((long)X & (sizeof (long) - 1)) | ((long)Y & (sizeof (long) - 1)))

/* How many bytes are copied each iteration of the 4X unrolled loop.  */
#define BIGBLOCKSIZE    (sizeof (long) << 2)

/* How many bytes are copied each iteration of the word copy loop.  */
#define LITTLEBLOCKSIZE (sizeof (long))

/* Threshold for punting to the byte copier.  */
#define TOO_SMALL(LEN)  ((LEN) < BIGBLOCKSIZE)

void *memcpy(void *dst0, const void *src0, size_t len0)
{
#if 0
    char *dstChar       = dst0;
    const char *srcChar = src0;
    for (size_t i = 0; i < len0; i++) {
        *(dstChar++) = *(srcChar++);
    }
    return dst0;
#else
    char *dst       = dst0;
    const char *src = src0;
    long *aligned_dst;
    const long *aligned_src;

    /* If the size is small, or either SRC or DST is unaligned,
       then punt into the byte copy loop.  This should be rare.  */
    if (!TOO_SMALL(len0) && !UNALIGNED(src, dst)) {
        aligned_dst = (long *)dst;
        aligned_src = (long *)src;

        /* Copy 4X long words at a time if possible.  */
        while (len0 >= BIGBLOCKSIZE) {
            *aligned_dst++ = *aligned_src++;
            *aligned_dst++ = *aligned_src++;
            *aligned_dst++ = *aligned_src++;
            *aligned_dst++ = *aligned_src++;
            len0 -= BIGBLOCKSIZE;
        }

        /* Copy one long word at a time if possible.  */
        while (len0 >= LITTLEBLOCKSIZE) {
            *aligned_dst++ = *aligned_src++;
            len0 -= LITTLEBLOCKSIZE;
        }

        /* Pick up any residual with a byte copier.  */
        dst = (char *)aligned_dst;
        src = (char *)aligned_src;
    }

    while (len0--)
        *dst++ = *src++;

    return dst0;
#endif
}

void *memmove(void *dst, const void *src, size_t n)
{
    char *dstChar       = dst;
    const char *srcChar = src;
    for (size_t i = 0; i < n; i++) {
        *(dstChar++) = *(srcChar++);
    }
    return dst;
}

void *memset(void *src, int c, size_t n)
{
    for (char *ptr = (char *)src; n > 0; n--, ptr++) {
        *ptr = (char)c;
    }
    return src;
}

char *itoa(long long int value, char *str, int base)
{
    char *rc;
    char *ptr;
    char *low;
    // Check for supported base.
    if (base < 2 || base > 36) {
        *str = '\0';
        return str;
    }
    rc = ptr = str;
    // Set '-' for negative decimals.
    if (value < 0 && base == 10) {
        *ptr++ = '-';
    }
    // Remember where the numbers start.
    low = ptr;
    // The actual conversion.
    do {
        // Modulo is negative for negative value. This trick makes abs() unnecessary.
        *ptr++ = "zyxwvutsrqponmlkjihgfedcba9876543210123456789abcdefghijklmnopqrstuvwxyz"
            [35 + value % base];
        value /= base;
    } while (value);
    // Terminating the string.
    *ptr-- = '\0';
    // Invert the numbers.
    while (low < ptr) {
        char tmp = *low;
        *low++   = *ptr;
        *ptr--   = tmp;
    }
    return rc;
}

int atoi(const char *str)
{
    int sign = 1, base = 0, i = 0;

    // if whitespaces then ignore.
    while (str[i] == ' ') {
        i++;
    }

    // sign of number
    if (str[i] == '-' || str[i] == '+') {
        sign = 1 - 2 * (str[i++] == '-');
    }

    // checking for valid input
    while (str[i] >= '0' && str[i] <= '9') {
        // handling overflow test case
        if (base > INT_MAX / 10 || (base == INT_MAX / 10 && str[i] - '0' > 7)) {
            if (sign == 1)
                return INT_MAX;
            else
                return INT_MIN;
        }
        base = 10 * base + (str[i++] - '0');
    }
    return base * sign;
}

/* K&R */
void reverse(char s[])
{
    int c, i, j;
    for (i = 0, j = strlen(s) - 1; i < j; i++, j--) {
        c    = s[i];
        s[i] = s[j];
        s[j] = c;
    }
}

/* K&R */
int strlen(const char s[])
{
    int i = 0;
    while (s[i] != '\0')
        ++i;
    return i;
}

/* K&R
 * Returns <0 if s1<s2, 0 if s1==s2, >0 if s1>s2 */
int strcmp(const char s1[], const char s2[])
{
    int i;
    for (i = 0; s1[i] == s2[i]; i++) {
        if (s1[i] == '\0')
            return 0;
    }
    return s1[i] - s2[i];
}

unsigned int strnlen(const char *s, size_t count)
{
    const char *sc;

    for (sc = s; count-- && *sc != '\0'; ++sc)
        /* nothing */ continue;

    return sc - s;
}

char *strzcpy(register char *dst, register const char *src, register int len)
{
    int i;

    if (len <= 0)
        return dst;

    for (i = 0; i < len; i++) {
        dst[i] = src[i];
        if (src[i] == '\0')
            return dst;
    }

    dst[len - 1] = '\0';
    return dst;
}

int printf(const char *format, ...)
{
    int ret;

    va_list ap;
    va_start(ap, format);
    ret = vprintf(format, ap);
    va_end(ap);

    return ret;
}

int puts(const char *str)
{
    int ret = 0;
    while (*str) {
        putc(*(str++));
        ret++;
    }
    return ret;
}

int putc(const char c)
{
    VGAPutc(c);
    serialPutc(c);
    return (unsigned char)c;
}

// int max is 2^(sizeof(int)*8) which is (2^3)^(sizeof(int)*8/3)
// = 8^(sizeof(int)*8/3) ~ 10^(sizeof(int)*8/3)
#define PRINT_INT(name, type)                                                                 \
    int print##name(type integer, char *str, size_t size)                                     \
    {                                                                                         \
        char num[sizeof(integer) * 3];                                                        \
        int i   = 0;                                                                          \
        int c   = 0;                                                                          \
        int ret = 0;                                                                          \
                                                                                              \
        if (integer < 0) {                                                                    \
            if (str) {                                                                        \
                if (size) {                                                                   \
                    str[c++] = '-';                                                           \
                    size--;                                                                   \
                    ret++;                                                                    \
                } else {                                                                      \
                    return ret;                                                               \
                }                                                                             \
            } else {                                                                          \
                ret++;                                                                        \
            }                                                                                 \
        }                                                                                     \
                                                                                              \
        do {                                                                                  \
            int digit = integer % 10;                                                         \
            num[i++]  = (digit > 0) ? digit : -digit;                                         \
            integer   = integer / 10;                                                         \
        } while (integer != 0);                                                               \
                                                                                              \
        for (i = i - 1; i >= 0; i--) {                                                        \
            if (str) {                                                                        \
                if (size) {                                                                   \
                    str[c++] = num[i] + '0';                                                  \
                    size--;                                                                   \
                    ret++;                                                                    \
                } else {                                                                      \
                    return ret;                                                               \
                }                                                                             \
            } else {                                                                          \
                ret++;                                                                        \
            }                                                                                 \
        }                                                                                     \
        return ret;                                                                           \
    }

#define PRINT_UINT(name, type)                                                                 \
    int print##name(type integer, char *str, size_t size)                                     \
    {                                                                                         \
        char num[sizeof(integer) * 3];                                                        \
        int i   = 0;                                                                          \
        int c   = 0;                                                                          \
        int ret = 0;                                                                          \
                                                                                              \
        do {                                                                                  \
            int digit = integer % 10;                                                         \
            num[i++]  = (digit > 0) ? digit : -digit;                                         \
            integer   = integer / 10;                                                         \
        } while (integer != 0);                                                               \
                                                                                              \
        for (i = i - 1; i >= 0; i--) {                                                        \
            if (str) {                                                                        \
                if (size) {                                                                   \
                    str[c++] = num[i] + '0';                                                  \
                    size--;                                                                   \
                    ret++;                                                                    \
                } else {                                                                      \
                    return ret;                                                               \
                }                                                                             \
            } else {                                                                          \
                ret++;                                                                        \
            }                                                                                 \
        }                                                                                     \
        return ret;                                                                           \
    }

PRINT_INT(Int, int);
PRINT_INT(Lint, long int);
PRINT_INT(Llint, long long int);
PRINT_UINT(Uint, unsigned int);
PRINT_UINT(Luint, long unsigned int);
PRINT_UINT(Lluint, long long unsigned int);

#define PRINT_PART(func, type, str, size, c, ret)                                             \
    {                                                                                         \
        int s;                                                                                \
        type d = va_arg(ap, type);                                                            \
        if (str)                                                                              \
            s = func(d, &str[c], size);                                                       \
        else                                                                                  \
            s = func(d, NULL, size);                                                          \
                                                                                              \
        size -= s;                                                                            \
        c += s;                                                                               \
        ret += s;                                                                             \
        break;                                                                                \
    }

int vsnprintf(char *str, size_t size, const char *format, va_list ap)
{
    int ret = 0;
    int i   = 0;
    int c   = 0;

    while (format[i] != '\0' && (size|| !str)) {
        switch (format[i]) {
            case '%':
                switch (format[i + 1]) {
                    case 'i':
                    case 'd': PRINT_PART(printInt, int, str, size, c, ret)
                    case 'u': PRINT_PART(printUint, uint, str, size, c, ret)
                    case 'p':
                    case 'x': {
                        char val[sizeof(int) * 2];
                        unsigned int valIdx = 0;
                        int d               = va_arg(ap, int);
                        itoa(d, val, 16);
                        if (str) {
                            while (val[valIdx]) {
                                if (size) {
                                    str[c++] = val[valIdx++];
                                    size--;
                                    ret++;
                                } else {
                                    return ret;
                                }
                            }
                        } else {
                            ret += strlen(val);
                        }
                        break;
                    }
                    case 'c': {
                        if (str) {
                            int ch   = va_arg(ap, int);
                            str[c++] = ch;
                            size--;
                        }
                        ret++;
                        break;
                    }
                    case 's': {
                        char *stri = va_arg(ap, char *);
                        if (!stri)
                            stri = "[NULL STR]";
                        if (str) {
                            while (*stri) {
                                if (size) {
                                    str[c++] = *(stri++);
                                    size--;
                                    ret++;
                                } else {
                                    return ret;
                                }
                            }
                        } else {
                            ret += strlen(stri);
                        }
                        break;
                    }
                    case '%':
                        if (str) {
                            str[c++] = '%';
                            size--;
                        }
                        ret++;
                        break;
                    case 'l':
                        switch (format[i + 2]) {
                            case 'l':
                                switch (format[i + 3]) {
                                    case 'i':
                                    case 'd': PRINT_PART(printLlint, long long int, str, size, c, ret)
                                    case 'u': PRINT_PART(printLluint, long long unsigned int, str, size, c, ret)
                                    case 'p':
                                    case 'x': {
                                        char val[sizeof(long long int) * 2];
                                        unsigned int valIdx = 0;
                                        unsigned long long int d =
                                            va_arg(ap, unsigned long long int);
                                        itoa(d, val, 16);
                                        if (str) {
                                            while (val[valIdx]) {
                                                if (size) {
                                                    str[c++] = val[valIdx++];
                                                    size--;
                                                    ret++;
                                                } else {
                                                    return ret;
                                                }
                                            }
                                        } else {
                                            ret += strlen(val);
                                        }
                                        break;
                                    }
                                }
                                i++;
                                break;
                            case 'i':
                            case 'd': PRINT_PART(printLint, long int, str, size, c, ret)
                            case 'u':
                                PRINT_PART(printLuint, long unsigned int, str, size, c, ret)
                            case 'p':
                            case 'x': {
                                char val[sizeof(int) * 2];
                                unsigned int valIdx = 0;
                                unsigned long int d = va_arg(ap, unsigned long int);
                                itoa(d, val, 16);
                                if (str) {
                                    while (val[valIdx]) {
                                        if (size) {
                                            str[c++] = val[valIdx++];
                                            size--;
                                            ret++;
                                        } else {
                                            return ret;
                                        }
                                    }
                                } else {
                                    ret += strlen(val);
                                }
                                break;
                            }
                        }
                        i++;
                        break;
                }
                i++;
                break;

            default:
                if (str) {
                    str[c++] = format[i];
                    size--;
                }
                ret++;
        }
        i++;
    }

    if (str) {
        if (size) {
            str[c++] = '\0';
        } else {
            if (c > 0) {
                str[c - 1] = '\0';
            }
        }
    }

    return ret;
}

int vprintf(const char *fmt, va_list ap)
{
    char tmp[256];
    int idx = 0;
    int ret = vsnprintf(tmp, sizeof(tmp), fmt, ap);
    while (ret > 0 && tmp[idx]) {
        putc(tmp[idx++]);
    }
    return ret;
}

int asprintf(char **strp, const char *fmt, ...)
{
    int ret;

    va_list ap;
    va_start(ap, fmt);
    ret = vasprintf(strp, fmt, ap);
    va_end(ap);

    return ret;
}

int vasprintf(char **strp, const char *fmt, va_list ap)
{
    int n       = 0;
    size_t size = 0;
    char *p     = NULL;

    /* Determine required size */

    n = vsnprintf(p, size, fmt, ap);

    if (n < 0)
        return -1;

    /* One extra byte for '\0' */

    size = (size_t)n + 1;
    p    = malloc(size);
    if (p == NULL)
        return -1;

    n = vsnprintf(p, size, fmt, ap);

    if (n < 0) {
        free(p);
        return -1;
    }
	*strp = p;

    return size;
}