#include <assert.h>
#include <string.h>
#include <stdlib.h>

#include "str.h"
#include "util.h"

#define INIT_SZ 20

// -----------------------------------------------------------
// INTERNAL: ensure capacity >= newsz (+1 for '\0')
// -----------------------------------------------------------
void
_str_ensure(str *s, int newsz)
{
    int tmp;

    newsz += 1;    // Account for null terminator

    if (newsz < s->capacity)
        return;

    tmp = s->capacity * 2;
    if (tmp > newsz)
        newsz = tmp;

    s->ptr = xrealloc(s->ptr, newsz);
    s->capacity = newsz;
}

// -----------------------------------------------------------
// Initialize empty string
// -----------------------------------------------------------
void
str_init(str *s)
{
    assert(s);

    s->capacity = INIT_SZ;
    s->len = 0;
    s->ptr = xmalloc(s->capacity);
    s->ptr[0] = '\0';
}

// -----------------------------------------------------------
// Free memory
// -----------------------------------------------------------
void
str_deinit(str *s)
{
    if (s->ptr)
        free(s->ptr);

    s->ptr = NULL;
}

// -----------------------------------------------------------
// Initialize with provided C-string
// -----------------------------------------------------------
void
str_init_data(str *s, const char *initial)
{
    assert(s && initial);

    int n = strlen(initial);
    s->capacity = (n + 1 < INIT_SZ) ? INIT_SZ : n + 1;
    s->len = n;

    s->ptr = xmalloc(s->capacity);
    memcpy(s->ptr, initial, n + 1); // copy including '\0'
}

// -----------------------------------------------------------
// Return pointer to underlying C-string
// -----------------------------------------------------------
char *
str_data(str *s)
{
    assert(s);
    return s->ptr;
}

// -----------------------------------------------------------
// Replace contents with provided C-string
// -----------------------------------------------------------
void
str_set(str *s, char *cstr)
{
    assert(s && cstr);

    int n = strlen(cstr);
    _str_ensure(s, n);

    memcpy(s->ptr, cstr, n + 1); // include '\0'
    s->len = n;
}

// -----------------------------------------------------------
// Create deep copy of string
// -----------------------------------------------------------
str
str_copy(str *s)
{
    assert(s);

    str s2;
    str_init(&s2);

    _str_ensure(&s2, s->len);
    memcpy(s2.ptr, s->ptr, s->len + 1);
    s2.len = s->len;

    return s2;
}

// -----------------------------------------------------------
// Count occurrences of character
// -----------------------------------------------------------
int
str_count(str *s, char ch)
{
    assert(s);

    int cnt = 0;
    for (int i = 0; i < s->len; i++)
        if (s->ptr[i] == ch)
            cnt++;

    return cnt;
}

// -----------------------------------------------------------
// Append C-string
// -----------------------------------------------------------
void
str_append(str *s, char *cstr)
{
    assert(s && cstr);

    int n = strlen(cstr);
    _str_ensure(s, s->len + n);

    memcpy(s->ptr + s->len, cstr, n + 1);
    s->len += n;
}

// -----------------------------------------------------------
// Append exactly len bytes from cstr
// -----------------------------------------------------------
void
str_append_n(str *s, char *cstr, int len)
{
    assert(s && cstr);

    _str_ensure(s, s->len + len);

    memcpy(s->ptr + s->len, cstr, len);
    s->len += len;

    s->ptr[s->len] = '\0'; // add terminator
}

// -----------------------------------------------------------
// Shrink string to new length
// -----------------------------------------------------------
void
str_shrink(str *s, int len)
{
    assert(s && len >= 0);

    if (len > s->len)
        len = s->len;

    s->len = len;
    s->ptr[len] = '\0';
}

// -----------------------------------------------------------
// Find substring, return index or -1
// -----------------------------------------------------------
int
str_find(str *s, char *substr)
{
    assert(s && substr);

    char *p = strstr(s->ptr, substr);
    if (!p)
        return -1;

    return (int)(p - s->ptr);
}

// -----------------------------------------------------------
// Extract substring [start_idx, start_idx+length)
// -----------------------------------------------------------
str
str_sub(str *s, int start_idx, int length)
{
    assert(s && start_idx >= 0 && length >= 0);

    str s2;
    str_init(&s2);

    if (start_idx > s->len)
        return s2; // empty

    if (start_idx + length > s->len)
        length = s->len - start_idx;

    _str_ensure(&s2, length);

    memcpy(s2.ptr, s->ptr + start_idx, length);
    s2.ptr[length] = '\0';
    s2.len = length;

    return s2;
}

// -----------------------------------------------------------
// Replace all occurrences of substr with replacement
// -----------------------------------------------------------
int
str_replace(str *s, char *substr, char *replacement)
{
    assert(s && substr && replacement);

    int sub_len = strlen(substr);
    int rep_len = strlen(replacement);

    if (sub_len == 0)
        return 0;

    str result;
    str_init(&result);

    char *pos = s->ptr;
    char *found;
    int count = 0;

    while ((found = strstr(pos, substr)) != NULL) {
        int prefix_len = (int)(found - pos);

        str_append_n(&result, pos, prefix_len);
        str_append(&result, replacement);

        pos = found + sub_len;
        count++;
    }

    if (*pos)
        str_append(&result, pos);

    str_set(s, result.ptr);
    str_deinit(&result);

    return count;
}