// retoor <retoor@molodetz.nl>
#include "tool.h"
#include "json_repair.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <pthread.h>
typedef struct {
    tool_t *tool;
    struct json_object *args;
    char *output;
    const char *name;
    const char *args_json;
} 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 *));
    bool *is_duplicate = calloc((size_t)len, sizeof(bool));
    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"));

        struct json_object *args_field = json_object_object_get(function_obj, "arguments");
        struct json_object *args = NULL;
        const char *args_json = NULL;

        if (args_field) {
            enum json_type atype = json_object_get_type(args_field);
            if (atype == json_type_object) {
                args = json_object_get(args_field);
                args_json = json_object_to_json_string(args);
            } else if (atype == json_type_string) {
                const char *raw = json_object_get_string(args_field);
                args_json = raw;
                if (raw) {
                    args = json_tokener_parse(raw);
                    if (!args) {
                        char *repaired = json_repair_string(raw);
                        if (repaired) {
                            args = json_tokener_parse(repaired);
                            free(repaired);
                        }
                    }
                }
            }
        }

        if (!args) {
            if (args_json) {
                size_t args_len = strlen(args_json);
                fprintf(stderr, "\033[1;31m[Registry] Failed to parse args for %s (len=%zu): %.200s...\033[0m\n",
                        name ? name : "unknown", args_len, args_json);
            }
            json_object_object_add(result_objs[i], "content",
                json_object_new_string(
                    "Error: failed to parse tool arguments (invalid JSON). "
                    "This is likely caused by output truncation — the arguments were cut off. "
                    "For large file content, use linux_terminal_execute with: "
                    "cat << 'HEREDOC_EOF' > filename\n...content...\nHEREDOC_EOF"));
            continue;
        }

        for (int j = 0; j < i; j++) {
            struct json_object *prev_call = json_object_array_get_idx(tool_calls, j);
            struct json_object *prev_func;
            json_object_object_get_ex(prev_call, "function", &prev_func);
            const char *prev_name = json_object_get_string(json_object_object_get(prev_func, "name"));
            const char *prev_args = t_args[j].args_json;
            if (name && prev_name && args_json && prev_args &&
                strcmp(name, prev_name) == 0 && strcmp(args_json, prev_args) == 0) {
                is_duplicate[i] = true;
                if (verbose) {
                    fprintf(stderr, "  \033[1;33m[Registry] Redundant call to %s prevented.\033[0m\n", name);
                }
                break;
            }
        }
        if (is_duplicate[i]) {
            json_object_put(args);
            continue;
        }
        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"));
            json_object_put(args);
            continue;
        }
        if (tool->vtable->print_action) {
            tool->vtable->print_action(tool->name, args);
        }
        if (verbose) {
            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;
        t_args[i].name = name;
        t_args[i].args_json = args_json;
        if (pthread_create(&threads[i], NULL, tool_thread_func, &t_args[i]) != 0) {
            threads[i] = 0;
            t_args[i].output = strdup("Error: failed to create thread for tool execution");
        }
    }
    for (int i = 0; i < len; i++) {
        if (!is_duplicate[i] && threads[i]) {
            pthread_join(threads[i], NULL);
        }
    }
    for (int i = 0; i < len; i++) {
        if (is_duplicate[i]) {
            struct json_object *curr_func;
            json_object_object_get_ex(json_object_array_get_idx(tool_calls, i), "function", &curr_func);
            const char *name = json_object_get_string(json_object_object_get(curr_func, "name"));
            const char *dup_args = t_args[i].args_json;

            for (int j = 0; j < i; j++) {
                const char *prev_name = t_args[j].name;
                const char *prev_args = t_args[j].args_json;
                if (name && prev_name && dup_args && prev_args &&
                    strcmp(name, prev_name) == 0 && strcmp(dup_args, prev_args) == 0) {
                    struct json_object *orig_content;
                    if (json_object_object_get_ex(result_objs[j], "content", &orig_content)) {
                        json_object_object_add(result_objs[i], "content", json_object_get(orig_content));
                    } else {
                        json_object_object_add(result_objs[i], "content", json_object_new_string("Result mirrored from previous parallel call."));
                    }
                    break;
                }
            }
        } else {
            char *output = t_args[i].output ? t_args[i].output : "";

            if (output[0] != '\0' && strncmp(output, "Error:", 6) == 0) {
                const char *tool_name = t_args[i].name ? t_args[i].name : "unknown";
                const char *received_args = t_args[i].args_json ? t_args[i].args_json : "{}";
                size_t enriched_len = strlen(output) + strlen(tool_name) + strlen(received_args) + 64;
                char *enriched = malloc(enriched_len);
                if (enriched) {
                    snprintf(enriched, enriched_len,
                             "Error in tool '%s': %s. Arguments received: %s",
                             tool_name, output + 7, received_args);
                    json_object_object_add(result_objs[i], "content", json_object_new_string(enriched));
                    fprintf(stderr, "\033[1;31m[Tool Error] %s\033[0m\n", enriched);
                    free(enriched);
                } else {
                    json_object_object_add(result_objs[i], "content", json_object_new_string(output));
                    fprintf(stderr, "\033[1;31m[Tool Error] %s\033[0m\n", output);
                }
            } else {
                json_object_object_add(result_objs[i], "content", json_object_new_string(output));
            }

            free(t_args[i].output);
            t_args[i].output = NULL;
            if (t_args[i].args) {
                json_object_put(t_args[i].args);
                t_args[i].args = NULL;
            }
        }
        json_object_array_add(results, result_objs[i]);
    }
    free(threads);
    free(t_args);
    free(result_objs);
    free(is_duplicate);
    return results;
}