// retoor <retoor@molodetz.nl>

#include "tool.h"
#include "agent.h"
#include "messages.h"
#include "r_config.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

typedef struct {
    tool_t tool;
} tool_agent_t;

static struct json_object *tool_spawn_agent_get_description(void) {
    struct json_object *obj = json_object_new_object();
    json_object_object_add(obj, "name", json_object_new_string("spawn_agent"));
    json_object_object_add(obj, "description", json_object_new_string("Spawn a specialized sub-agent to handle a specific task."));

    struct json_object *params = json_object_new_object();
    json_object_object_add(params, "type", json_object_new_string("object"));

    struct json_object *props = json_object_new_object();

    struct json_object *persona = json_object_new_object();
    json_object_object_add(persona, "type", json_object_new_string("string"));
    json_object_object_add(persona, "description", json_object_new_string("The persona of the agent (researcher, developer, security)."));
    struct json_object *persona_enum = json_object_new_array();
    json_object_array_add(persona_enum, json_object_new_string("researcher"));
    json_object_array_add(persona_enum, json_object_new_string("developer"));
    json_object_array_add(persona_enum, json_object_new_string("security"));
    json_object_object_add(persona, "enum", persona_enum);
    json_object_object_add(props, "persona", persona);

    struct json_object *goal = json_object_new_object();
    json_object_object_add(goal, "type", json_object_new_string("string"));
    json_object_object_add(goal, "description", json_object_new_string("The specific task or goal for the sub-agent."));
    json_object_object_add(props, "goal", goal);

    json_object_object_add(params, "properties", props);

    struct json_object *required = json_object_new_array();
    json_object_array_add(required, json_object_new_string("persona"));
    json_object_array_add(required, json_object_new_string("goal"));
    json_object_object_add(params, "required", required);

    json_object_object_add(obj, "parameters", params);

    struct json_object *full_obj = json_object_new_object();
    json_object_object_add(full_obj, "type", json_object_new_string("function"));
    json_object_object_add(full_obj, "function", obj);

    return full_obj;
}

static char *tool_spawn_agent_execute(tool_t *self, struct json_object *args) {
    (void)self;
    struct json_object *persona_obj, *goal_obj;

    if (!json_object_object_get_ex(args, "persona", &persona_obj) ||
        !json_object_object_get_ex(args, "goal", &goal_obj)) {
        return strdup("Error: Missing persona or goal");
    }

    const char *persona_str = json_object_get_string(persona_obj);
    const char *goal_str = json_object_get_string(goal_obj);

    tool_registry_type_t type = TOOL_TYPE_ALL;
    const char *system_prompt = NULL;

    if (strcmp(persona_str, "researcher") == 0) {
        type = TOOL_TYPE_RESEARCHER;
        system_prompt = "You are a specialized Research Agent. Your goal is to find, extract, and summarize information. "
                        "Do not attempt to write or execute code unless it's for data analysis. "
                        "Focus on using web search, http fetch, and reading files.";
    } else if (strcmp(persona_str, "developer") == 0) {
        type = TOOL_TYPE_DEVELOPER;
        system_prompt = "You are a specialized Developer Agent. Your goal is to write, test, and debug code. "
                        "Use the terminal, file editing tools, and python execution to fulfill your task. "
                        "Always verify your changes by running tests or the code itself.";
    } else if (strcmp(persona_str, "security") == 0) {
        type = TOOL_TYPE_SECURITY;
        system_prompt = "You are a specialized Security Auditor Agent. Your goal is to find vulnerabilities and perform security analysis. "
                        "Be pedantic and thorough. Use fuzzing, port scanning, and code analysis tools. "
                        "Report any findings clearly with potential impact.";
    } else {
        return strdup("Error: Invalid persona");
    }

    char session_id[256];
    snprintf(session_id, sizeof(session_id), "subagent-%s-%u", persona_str, (unsigned int)time(NULL));

    messages_handle msgs = messages_create(session_id);
    messages_add(msgs, "system", system_prompt);

    agent_handle agent = agent_create(goal_str, msgs);
    if (!agent) {
        messages_destroy(msgs);
        return strdup("Error: Failed to create sub-agent");
    }

    tool_registry_t *specialized_tools = tool_registry_get_specialized(type);
    agent_set_tool_registry(agent, specialized_tools);
    agent_set_max_iterations(agent, 50); // Sub-agents have lower limit

    char *result = agent_run(agent, goal_str);

    agent_destroy(agent);
    // specialized_tools should be destroyed? 
    // In tool_registry_get_specialized we create a new one.
    tool_registry_destroy(specialized_tools);

    if (!result) {
        return strdup("Sub-agent failed to provide a result.");
    }

    return result;
}

static void tool_spawn_agent_print_action(const char *name, struct json_object *args) {
    struct json_object *persona_obj, *goal_obj;
    const char *persona = "unknown";
    const char *goal = "unknown";

    if (json_object_object_get_ex(args, "persona", &persona_obj)) persona = json_object_get_string(persona_obj);
    if (json_object_object_get_ex(args, "goal", &goal_obj)) goal = json_object_get_string(goal_obj);

    printf("\033[1;34m[Agent] Spawning %s agent for: %s\033[0m\n", persona, goal);
}

static const tool_vtable_t tool_spawn_agent_vtable = {
    .get_description = tool_spawn_agent_get_description,
    .execute = tool_spawn_agent_execute,
    .print_action = tool_spawn_agent_print_action
};

tool_t *tool_spawn_agent_create(void) {
    tool_agent_t *tool = calloc(1, sizeof(tool_agent_t));
    if (!tool) return NULL;

    tool->tool.vtable = &tool_spawn_agent_vtable;
    tool->tool.name = "spawn_agent";

    return &tool->tool;
}