#include "str.h"

#include <assert.h>

#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>

static inline bool ystr_is_inline(ystr_t *s) { return 0 == s->inner.big.size; }

ystr_t ystr_new(const char* src) {
    ystr_t str;
    ystr_init(&str, src);
    return str;
}

void ystr_init(ystr_t *s, const char *src) {
    memset(s, 0, sizeof(ystr_t));
    if (src == NULL) return;
    ystr_set(s, src);
}

void ystr_set_owned(ystr_t *s, char *src) {
    if (src == NULL) {
        ystr_destroy(s);
        return;
    }
    // Free previous memory.
    ystr_destroy(s);
    size_t len = strlen(src) + 1;
    s->inner.big.len = len;
    s->inner.big.cap = len;
    s->inner.big.size = 1;
    s->inner.big.base = src;
}

void ystr_set(ystr_t *s, const char *src) {
    if (src == NULL) {
        ystr_destroy(s);
        return;
    }
    ystr_set_n(s, src, strlen(src));
}

void ystr_set_n(ystr_t *s, const char *src, size_t srclen) {
    if (src == NULL) {
        ystr_destroy(s);
        return;
    }
    if (srclen < YSTR_SMALL_THRESHOLD) {
        if (!ystr_is_inline(s))
            yvec_destroy(&s->inner.big);
        memset(s, 0,
                sizeof(ystr_t)); // Marks as inline because inner.big.size == 0.
        memcpy(s->inner.small, src, srclen);
    } else {
        // +1 for terminating NULL byte.
        if (ystr_is_inline(s))
            yvec_init(&s->inner.big, sizeof(char), srclen + 1);
        yvec_resize(&s->inner.big, srclen + 1);
        memcpy(s->inner.big.base, src, srclen + 1);
    }
}

void ystr_destroy(ystr_t *s) {
    if (!ystr_is_inline(s)) {
        free(s->inner.big.base);
    }
    memset(s, 0, sizeof(ystr_t));
}

size_t ystr_len(ystr_t *s) {
    if (ystr_is_inline(s))
        return strlen(s->inner.small);
    return s->inner.big.len - 1;
}

char* ystr_at(ystr_t *s, size_t index) {
    if (index >= ystr_len(s))
        return NULL;
    if (ystr_is_inline(s)) {
        return &s->inner.small[index];
    }
    return YVEC_AT(&s->inner.big, index, char);
}

const char *ystr_str(ystr_t *s) {
    if (ystr_is_inline(s)) {
        return s->inner.small;
    }
    return (const char *)s->inner.big.base;
}

int ystr_cmp(ystr_t *s1, ystr_t *s2) {
    return strcmp(ystr_str(s1), ystr_str(s2));
}

int ystr_cmp_str(ystr_t *s1, const char *s2) {
    return strcmp(ystr_str(s1), s2);
}

void ystr_split(ystr_t *s, char limit, yvec_t *dst) {
    size_t parts = 0;
    size_t len = ystr_len(s);
    for (size_t i = 0; i < len; i++) {
        if (*ystr_at(s, i) == limit)
            parts++;
    }
    parts += 1;
    yvec_init(dst, sizeof(ystr_t), parts);
    size_t index = 0;
    while (index <= len) {
        ystr_t part;
        ystr_init(&part, NULL);
        size_t i = 0;
        for (i = index; i < len && *ystr_at(s, i) != limit; i++)
            ;
        ystr_set_n(&part, ystr_str(s) + index, i - index);
        YVEC_PUSH(dst, &part);
        index = i + 1;
    }
}

void ystr_resize(ystr_t *s, size_t new) {
    if (ystr_is_inline(s)) {
        size_t oldlen = strlen(s->inner.small);
        if (new < oldlen) {
            // Truncate small string.
            memset(s->inner.small + new, 0, YSTR_SMALL_THRESHOLD - new);
        } else if (new > YSTR_SMALL_THRESHOLD) {
            // Move to external representation.
            yvec_t ext;
            YVEC_INIT(&ext, new + 1, char);
            yvec_push_multi(&ext, s->inner.small, oldlen + 1);
            yvec_resize(&ext, new + 1);
            s->inner.big = ext;
        } else {
            // Size is between current size and threshold; do nothing.
        }
        return;
    }

    // +1 for NULL byte.
    yvec_resize(&s->inner.big, new + 1);
}

void ystr_append(ystr_t *s1, ystr_t *s2) {
    size_t oldlen = ystr_len(s1);
    ystr_resize(s1, ystr_len(s1) + ystr_len(s2));

    // Still inline after resize?
    if (ystr_is_inline(s1)) {
        memcpy(s1->inner.small + strlen(s1->inner.small), ystr_str(s2),
                ystr_len(s2));
        return;
    }
    // String has external vec with length of new string. Resize to length of
    // old string without NULL before appending to it (append starts at current
    // length, we want to skip the NULL byte)
    yvec_resize(&s1->inner.big, oldlen);
    yvec_push_multi(&s1->inner.big, ystr_str(s2), ystr_len(s2) + 1);
}

void ystr_copy(ystr_t* src, ystr_t* dst) {
    if (dst == NULL || src == NULL)
        return;
    ystr_destroy(dst);
    if (ystr_is_inline(src)) {
        memcpy(dst, src, sizeof(ystr_t));
        return;
    }
    yvec_copy(&src->inner.big, &dst->inner.big);
}

ystr_t ystr_clone(ystr_t* src) {
    ystr_t new;
    ystr_init(&new, ystr_str(src));
    return new;
}

#define YSTR_BUILD_BUF_SIZE 128

void ystr_build(ystr_t *s, const char *fmt, ...) {
    va_list args;
    char buf[YSTR_BUILD_BUF_SIZE];
    memset(buf, 0, YSTR_BUILD_BUF_SIZE);

    va_start(args, fmt);
    int written = vsnprintf(buf, YSTR_BUILD_BUF_SIZE, fmt, args);
    va_end(args);

    // If truncated, make more space.
    if (written >= YSTR_BUILD_BUF_SIZE) {
        char *buf2 = calloc(written + 1, sizeof(char));
        va_start(args, fmt);
        vsnprintf(buf2, written + 1, fmt, args);
        va_end(args);

        ystr_t fmtd;
        ystr_init(&fmtd, NULL);
        ystr_set_owned(&fmtd, buf2);
        ystr_append(s, &fmtd);
        ystr_destroy(&fmtd);
        return;
    }
    ystr_t fmtd;
    ystr_init(&fmtd, NULL);
    ystr_set_owned(&fmtd, buf);
    ystr_append(s, &fmtd);
    // No destroy, fmtd points to buffer on stack.
}