// retoor <retoor@molodetz.nl>

#include "tool.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <pthread.h>

typedef struct {
    tool_t *tool;
    struct json_object *args;
    char *output;
} tool_thread_args_t;

static void *tool_thread_func(void *ptr) {
    tool_thread_args_t *args = (tool_thread_args_t *)ptr;
    if (args->tool->vtable->execute) {
        args->output = args->tool->vtable->execute(args->tool, args->args);
    }
    return NULL;
}

tool_registry_t *tool_registry_create(void) {
    tool_registry_t *registry = calloc(1, sizeof(tool_registry_t));
    if (!registry) return NULL;

    registry->capacity = 32;
    registry->tools = calloc(registry->capacity, sizeof(tool_t *));
    if (!registry->tools) {
        free(registry);
        return NULL;
    }
    return registry;
}

void tool_registry_destroy(tool_registry_t *registry) {
    if (!registry) return;
    free(registry->tools);
    free(registry);
}

r_status_t tool_registry_register(tool_registry_t *registry, tool_t *tool) {
    if (!registry || !tool) return R_ERROR_INVALID_ARG;

    if (registry->count >= registry->capacity) {
        size_t new_capacity = registry->capacity * 2;
        tool_t **new_tools = realloc(registry->tools, new_capacity * sizeof(tool_t *));
        if (!new_tools) return R_ERROR_OUT_OF_MEMORY;
        registry->tools = new_tools;
        registry->capacity = new_capacity;
    }

    registry->tools[registry->count++] = tool;
    return R_SUCCESS;
}

tool_t *tool_registry_find(tool_registry_t *registry, const char *name) {
    if (!registry || !name) return NULL;

    for (size_t i = 0; i < registry->count; i++) {
        if (strcmp(registry->tools[i]->name, name) == 0) {
            return registry->tools[i];
        }
    }
    return NULL;
}

struct json_object *tool_registry_get_descriptions(tool_registry_t *registry) {
    if (!registry) return NULL;

    struct json_object *array = json_object_new_array();
    if (!array) return NULL;

    for (size_t i = 0; i < registry->count; i++) {
        tool_t *tool = registry->tools[i];
        if (tool->vtable->get_description) {
            struct json_object *desc = tool->vtable->get_description();
            if (desc) {
                json_object_array_add(array, desc);
            }
        }
    }
    return array;
}

struct json_object *tool_registry_execute(tool_registry_t *registry,
                                          struct json_object *tool_calls,
                                          bool verbose) {
    if (!registry || !tool_calls) return NULL;

    struct json_object *results = json_object_new_array();
    if (!results) return NULL;

    int len = json_object_array_length(tool_calls);
    if (len == 0) return results;

    pthread_t *threads = calloc((size_t)len, sizeof(pthread_t));
    tool_thread_args_t *t_args = calloc((size_t)len, sizeof(tool_thread_args_t));
    struct json_object **result_objs = calloc((size_t)len, sizeof(struct json_object *));

    for (int i = 0; i < len; i++) {
        struct json_object *call = json_object_array_get_idx(tool_calls, i);
        result_objs[i] = json_object_new_object();

        struct json_object *id_obj = json_object_object_get(call, "id");
        if (id_obj) {
            json_object_object_add(result_objs[i], "tool_call_id",
                                   json_object_new_string(json_object_get_string(id_obj)));
        }
        json_object_object_add(result_objs[i], "role", json_object_new_string("tool"));

        struct json_object *function_obj;
        if (!json_object_object_get_ex(call, "function", &function_obj)) {
            json_object_object_add(result_objs[i], "content", json_object_new_string("Error: missing function"));
            continue;
        }

        const char *name = json_object_get_string(json_object_object_get(function_obj, "name"));
        tool_t *tool = tool_registry_find(registry, name);

        if (!tool) {
            json_object_object_add(result_objs[i], "content", json_object_new_string("Error: tool not found"));
            continue;
        }

        struct json_object *args = NULL;
        struct json_object *args_obj;
        if (json_object_object_get_ex(function_obj, "arguments", &args_obj)) {
            args = json_tokener_parse(json_object_get_string(args_obj));
        }

        if (tool->vtable->print_action) {
            tool->vtable->print_action(tool->name, args);
        }

        if (verbose && args) {
            fprintf(stderr, "  \033[2m[parallel] launching %s\033[0m\n", tool->name);
        }

        t_args[i].tool = tool;
        t_args[i].args = args;
        t_args[i].output = NULL;

        pthread_create(&threads[i], NULL, tool_thread_func, &t_args[i]);
    }

    for (int i = 0; i < len; i++) {
        if (threads[i]) {
            pthread_join(threads[i], NULL);
        }
        json_object_object_add(result_objs[i], "content",
                               json_object_new_string(t_args[i].output ? t_args[i].output : ""));
        free(t_args[i].output);
        if (t_args[i].args) json_object_put(t_args[i].args);
        json_object_array_add(results, result_objs[i]);
    }

    free(threads);
    free(t_args);
    free(result_objs);

    return results;
}