#include <bits/errno-return.h>
#include <errno.h>
#include <luna/HashTable.h>
#include <luna/ScopeGuard.h>
#include <luna/Vector.h>
#include <stdlib.h>

extern "C" char** environ;
Vector<char*> g_dynamic_env = {};
HashTable<char*> g_dynamic_vars = {};
bool env_is_dynamic { false };

/* If an environment variable matching key is found, stores its value in value and returns its index. Otherwise,
 * returns -1. */
static isize _findenv(const char* key, char** value)
{
    char** env = environ;
    size_t len = strlen(key);

    if (!env) return -1;

    while (*env)
    {
        // Retrieve an environment variable from environ.
        char* var = *(env++);

        // Check for an equals sign, else we just skip this one.
        char* delim = strchr(var, '=');
        if (!delim) continue;

        // Get the length of this variable's key (the part before the equals sign)
        size_t key_length = strcspn(var, "=");
        if (len != key_length) continue;

        // If the keys match, we found it! Return the value (which is just after the equals sign)
        if (!memcmp(key, var, key_length))
        {
            if (value) *value = delim + 1;
            return (env - environ) - 1;
        }
    }

    return -1;
}

static void free_if_needed(char* element)
{
    if (g_dynamic_vars.try_find(element))
    {
        g_dynamic_vars.try_remove(element);
        free(element);
    }
}

/* Move environ to a heap-allocated array. */
static Result<void> _try_move_env()
{
    char** env = environ;
    g_dynamic_env.clear();

    if (!env)
    {
        TRY(g_dynamic_env.try_append(nullptr));
        env_is_dynamic = true;
        environ = g_dynamic_env.data();
        return {};
    }

    while (*env)
    {
        char* ptr = *(env++);
        TRY(g_dynamic_env.try_append(ptr));
    }

    TRY(g_dynamic_env.try_append(nullptr));

    env_is_dynamic = true;
    environ = g_dynamic_env.data();

    return {};
}

static int _move_env()
{
    TRY_OR_SET_ERRNO(_try_move_env(), int, -1);

    return 0;
}

static void _update_env()
{
    environ = g_dynamic_env.data();
}

// Check if a user has modified environ, in which case we should reset env_is_dynamic and free our buffers.
static void _check_dynamic_env()
{
    if (!env_is_dynamic) return;

    env_is_dynamic = environ == g_dynamic_env.data();

    if (!env_is_dynamic)
    {
        for (auto element : g_dynamic_env)
        {
            if (element) free_if_needed(element);
        }
    }
}

extern "C"
{
    char** environ = nullptr;

    int clearenv()
    {
        _check_dynamic_env();

        if (env_is_dynamic)
        {
            for (auto element : g_dynamic_env)
            {
                if (element) free_if_needed(element);
            }

            g_dynamic_env.clear();
        }

        environ = nullptr;
        env_is_dynamic = false;

        return 0;
    }

    int unsetenv(const char* key)
    {
        if (!key || *key == 0 || strchr(key, '='))
        {
            errno = EINVAL;
            return -1;
        }

        auto index = _findenv(key, nullptr);
        if (index < 0) return 0;

        _check_dynamic_env();

        if (!env_is_dynamic)
        {
            if (_move_env() < 0) return -1;
        }

        char* p = g_dynamic_env.remove_at(index);
        _update_env();

        free_if_needed(p);

        return 0;
    }

    int setenv(const char* key, const char* value, int overwrite)
    {
        if (!key || *key == 0 || strchr(key, '='))
        {
            errno = EINVAL;
            return -1;
        }

        auto index = _findenv(key, nullptr);
        if (index >= 0 && !overwrite) return 0;

        _check_dynamic_env();

        if (!env_is_dynamic)
        {
            if (_move_env() < 0) return -1;
        }

        usize len = strlen(key) + strlen(value) + 2;

        char* str = (char*)calloc(len, 1);
        if (!str) return -1;
        auto guard = make_scope_guard([=] { free(str); });

        strncpy(str, key, strlen(key) + 1);
        strncat(str, "=", 2);
        strncat(str, value, strlen(value) + 1);

        if (index >= 0)
        {
            free_if_needed(environ[index]);
            environ[index] = str;
            guard.deactivate();
            return 0;
        }

        TRY_OR_SET_ERRNO(g_dynamic_vars.try_set(str), int, -1);

        // Add a new NULL at the end of the array and replace the previous one with our string.
        index = g_dynamic_env.size() - 1;

        TRY_OR_SET_ERRNO(g_dynamic_env.try_append(nullptr), int, -1);
        guard.deactivate();
        _update_env();

        environ[index] = str;

        return 0;
    }

    int putenv(char* string)
    {
        char* p = strchr(string, '=');
        if (!p) return unsetenv(string);
        size_t key_len = p - string;

        char* key = strndup(string, key_len);
        auto guard = make_scope_guard([key] { free(key); });

        auto index = _findenv(key, nullptr);

        _check_dynamic_env();

        if (!env_is_dynamic)
        {
            if (_move_env() < 0) return -1;
        }

        if (index >= 0)
        {
            free_if_needed(environ[index]);
            environ[index] = string;
            return 0;
        }

        // Add a new NULL at the end of the array and replace the previous one with our string.
        index = g_dynamic_env.size() - 1;

        TRY_OR_SET_ERRNO(g_dynamic_env.try_append(nullptr), int, -1);
        guard.deactivate();
        _update_env();

        environ[index] = string;

        return 0;
    }

    char* getenv(const char* key)
    {
        char* result;
        if (_findenv(key, &result) < 0) return nullptr;
        return result;
    }
}