// retoor <retoor@molodetz.nl>

#include "agent.h"
#include "http_client.h"
#include "r_config.h"
#include "tool.h"
#include <json-c/json.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

struct agent_t {
    char *goal;
    int iteration_count;
    int max_iterations;
    int tool_retry_count;
    int max_tool_retries;
    agent_state_t state;
    time_t start_time;
    char *last_error;
    bool verbose;
    messages_handle messages;
    bool owns_messages;
    http_client_handle http;
    tool_registry_t *tools;
};

static const char *incomplete_phrases[] = {
    "I'll ",      "I will ",    "Let me ",     "I'm going to ",
    "Next, I",    "Now I'll",   "Now I will",  "I'll now",
    "I need to",  "I should",   "I can ",      "Going to ",
    "Will now",   "Proceeding", "Starting to", "About to ",
    "First, I",   "Then I",     "After that",  "Following that",
    NULL
};

static const char *incomplete_endings[] = {
    "...", ":", "files:", "content:", "implementation:",
    NULL
};

extern tool_registry_t *tools_get_registry(void);

agent_handle agent_create(const char *goal, messages_handle messages) {
    struct agent_t *agent = calloc(1, sizeof(struct agent_t));
    if (!agent) return NULL;

    if (goal) {
        agent->goal = strdup(goal);
        if (!agent->goal) {
            free(agent);
            return NULL;
        }
    }

    r_config_handle cfg = r_config_get_instance();

    agent->iteration_count = 0;
    agent->max_iterations = AGENT_MAX_ITERATIONS;
    agent->tool_retry_count = 0;
    agent->max_tool_retries = AGENT_MAX_TOOL_RETRIES;
    agent->state = AGENT_STATE_IDLE;
    agent->start_time = time(NULL);
    agent->verbose = r_config_is_verbose(cfg);

    if (messages) {
        agent->messages = messages;
        agent->owns_messages = false;
    } else {
        agent->messages = messages_create(r_config_get_session_id(cfg));
        agent->owns_messages = true;
    }

    if (!agent->messages) {
        free(agent->goal);
        free(agent);
        return NULL;
    }

    agent->http = http_client_create(r_config_get_api_key(cfg));
    if (!agent->http) {
        if (agent->owns_messages) {
            messages_destroy(agent->messages);
        }
        free(agent->goal);
        free(agent);
        return NULL;
    }

    agent->tools = tools_get_registry();

    return agent;
}

void agent_destroy(agent_handle agent) {
    if (!agent) return;
    if (agent->http) http_client_destroy(agent->http);
    if (agent->messages && agent->owns_messages) messages_destroy(agent->messages);
    free(agent->goal);
    free(agent->last_error);
    free(agent);
}

void agent_set_max_iterations(agent_handle agent, int max) {
    if (agent) agent->max_iterations = max;
}

void agent_set_verbose(agent_handle agent, bool verbose) {
    if (agent) agent->verbose = verbose;
}

agent_state_t agent_get_state(agent_handle agent) {
    return agent ? agent->state : AGENT_STATE_ERROR;
}

const char *agent_get_error(agent_handle agent) {
    return agent ? agent->last_error : NULL;
}

int agent_get_iteration_count(agent_handle agent) {
    return agent ? agent->iteration_count : 0;
}

static void agent_set_error(agent_handle agent, const char *error) {
    if (!agent) return;
    free(agent->last_error);
    agent->last_error = error ? strdup(error) : NULL;
}

static char *agent_build_request(agent_handle agent, const char *role, const char *message) {
    r_config_handle cfg = r_config_get_instance();

    struct json_object *root = json_object_new_object();
    if (!root) return NULL;

    json_object_object_add(root, "model",
        json_object_new_string(r_config_get_model(cfg)));

    if (role && message) {
        messages_add(agent->messages, role, message);
        if (r_config_use_tools(cfg) && agent->tools) {
            json_object_object_add(root, "tools",
                tool_registry_get_descriptions(agent->tools));
        }
    }

    json_object_object_add(root, "messages",
        json_object_get(messages_to_json(agent->messages)));
    json_object_object_add(root, "temperature",
        json_object_new_double(r_config_get_temperature(cfg)));

    char *result = strdup(json_object_to_json_string_ext(root, JSON_C_TO_STRING_PRETTY));
    json_object_put(root);
    return result;
}

static struct json_object *agent_process_response(agent_handle agent, const char *json_data) {
    r_config_handle cfg = r_config_get_instance();

    char *response = NULL;
    r_status_t status = http_post(agent->http, r_config_get_api_url(cfg), json_data, &response);

    if (status != R_SUCCESS || !response) {
        return NULL;
    }

    struct json_object *parsed = json_tokener_parse(response);
    free(response);

    if (!parsed) return NULL;

    struct json_object *error_obj;
    if (json_object_object_get_ex(parsed, "error", &error_obj)) {
        const char *err_str = json_object_to_json_string(error_obj);
        fprintf(stderr, "API Error: %s\n", err_str);
        json_object_put(parsed);
        return NULL;
    }

    struct json_object *choices;
    if (!json_object_object_get_ex(parsed, "choices", &choices)) {
        json_object_put(parsed);
        return NULL;
    }

    struct json_object *first_choice = json_object_array_get_idx(choices, 0);
    if (!first_choice) {
        json_object_put(parsed);
        return NULL;
    }

    return first_choice;
}

static bool agent_has_tool_calls(struct json_object *choice) {
    struct json_object *message_obj;
    if (!json_object_object_get_ex(choice, "message", &message_obj)) return false;

    struct json_object *tool_calls;
    if (!json_object_object_get_ex(message_obj, "tool_calls", &tool_calls)) return false;

    return json_object_array_length(tool_calls) > 0;
}

static struct json_object *agent_get_tool_calls(struct json_object *choice) {
    struct json_object *message_obj;
    if (!json_object_object_get_ex(choice, "message", &message_obj)) return NULL;

    struct json_object *tool_calls;
    if (!json_object_object_get_ex(message_obj, "tool_calls", &tool_calls)) return NULL;

    return tool_calls;
}

static struct json_object *agent_get_message(struct json_object *choice) {
    struct json_object *message_obj;
    if (json_object_object_get_ex(choice, "message", &message_obj)) {
        return message_obj;
    }
    return NULL;
}

static char *agent_get_content(struct json_object *choice) {
    struct json_object *message_obj;
    if (!json_object_object_get_ex(choice, "message", &message_obj)) return NULL;

    struct json_object *content_obj;
    if (!json_object_object_get_ex(message_obj, "content", &content_obj)) return NULL;

    const char *content = json_object_get_string(content_obj);
    return content ? strdup(content) : NULL;
}

static bool agent_response_indicates_incomplete(const char *content) {
    if (!content) return false;

    for (int i = 0; incomplete_phrases[i]; i++) {
        if (strstr(content, incomplete_phrases[i])) return true;
    }

    size_t len = strlen(content);
    if (len > 3) {
        for (int i = 0; incomplete_endings[i]; i++) {
            size_t end_len = strlen(incomplete_endings[i]);
            if (len >= end_len && strcmp(content + len - end_len, incomplete_endings[i]) == 0) {
                return true;
            }
        }
    }

    return false;
}

char *agent_run(agent_handle agent, const char *user_message) {
    if (!agent) return NULL;

    agent->state = AGENT_STATE_RUNNING;
    agent->iteration_count = 0;
    agent->tool_retry_count = 0;

    if (!user_message || !*user_message) {
        agent->state = AGENT_STATE_ERROR;
        agent_set_error(agent, "Empty user message");
        return NULL;
    }

    messages_load(agent->messages);

    char *json_data = agent_build_request(agent, "user", user_message);
    if (!json_data) {
        agent->state = AGENT_STATE_ERROR;
        agent_set_error(agent, "Failed to create chat JSON");
        return NULL;
    }

    char *final_response = NULL;

    while (agent->state == AGENT_STATE_RUNNING) {
        agent->iteration_count++;

        if (agent->iteration_count > agent->max_iterations) {
            agent->state = AGENT_STATE_MAX_ITERATIONS;
            agent_set_error(agent, "Maximum iterations reached");
            if (agent->verbose) {
                fprintf(stderr, "[Agent] Max iterations (%d) reached\n", agent->max_iterations);
            }
            free(json_data);
            break;
        }

        if (agent->verbose) {
            fprintf(stderr, "[Agent] Iteration %d/%d\n",
                agent->iteration_count, agent->max_iterations);
        }

        struct json_object *choice = agent_process_response(agent, json_data);
        free(json_data);
        json_data = NULL;

        if (!choice) {
            agent->tool_retry_count++;
            if (agent->tool_retry_count >= agent->max_tool_retries) {
                agent->state = AGENT_STATE_ERROR;
                agent_set_error(agent, "API request failed after retries");
                break;
            }
            if (agent->verbose) {
                fprintf(stderr, "[Agent] API error, retry %d/%d\n",
                    agent->tool_retry_count, agent->max_tool_retries);
            }
            json_data = agent_build_request(agent, NULL, NULL);
            continue;
        }

        agent->tool_retry_count = 0;

        struct json_object *message_obj = agent_get_message(choice);
        if (message_obj) {
            messages_add_object(agent->messages, json_object_get(message_obj));
        }

        bool has_tools = agent_has_tool_calls(choice);

        if (agent->verbose) {
            fprintf(stderr, "[Agent] has_tool_calls=%s\n", has_tools ? "true" : "false");
        }

        if (has_tools) {
            agent->state = AGENT_STATE_EXECUTING_TOOLS;

            struct json_object *tool_calls = agent_get_tool_calls(choice);

            if (agent->verbose) {
                int num_tools = json_object_array_length(tool_calls);
                fprintf(stderr, "[Agent] Executing %d tool(s)\n", num_tools);
            }

            struct json_object *results = tool_registry_execute(agent->tools, tool_calls, agent->verbose);

            int count = json_object_array_length(results);
            for (int i = 0; i < count; i++) {
                struct json_object *result = json_object_array_get_idx(results, i);
                messages_add_tool_call(agent->messages, json_object_get(result));
            }

            agent->state = AGENT_STATE_RUNNING;
            json_data = agent_build_request(agent, NULL, NULL);
            if (!json_data) {
                agent->state = AGENT_STATE_ERROR;
                agent_set_error(agent, "Failed to create follow-up JSON");
                break;
            }

        } else {
            char *content = agent_get_content(choice);

            if (content && agent_response_indicates_incomplete(content)) {
                if (agent->verbose) {
                    fprintf(stderr, "[Agent] Response indicates incomplete work, auto-continuing\n");
                }

                free(content);
                json_data = agent_build_request(agent, "user",
                    "Continue. Execute the necessary actions to complete the task.");
                if (!json_data) {
                    agent->state = AGENT_STATE_ERROR;
                    agent_set_error(agent, "Failed to create continue JSON");
                    break;
                }

            } else {
                final_response = content;
                agent->state = AGENT_STATE_COMPLETED;

                if (agent->verbose) {
                    fprintf(stderr, "[Agent] Completed in %d iteration(s)\n",
                        agent->iteration_count);
                }
            }
        }
    }

    free(json_data);
    return final_response;
}

char *agent_chat(const char *user_message, messages_handle messages) {
    agent_handle agent = agent_create(user_message, messages);
    if (!agent) return NULL;

    char *response = agent_run(agent, user_message);

    if (agent->verbose && agent->state != AGENT_STATE_COMPLETED && agent->last_error) {
        fprintf(stderr, "[Agent] Error: %s\n", agent->last_error);
    }

    agent_destroy(agent);
    return response;
}

char *agent_chat_with_limit(const char *user_message, int max_iterations, messages_handle messages) {
    agent_handle agent = agent_create(user_message, messages);
    if (!agent) return NULL;

    agent_set_max_iterations(agent, max_iterations);
    char *response = agent_run(agent, user_message);

    if (agent->verbose && agent->state != AGENT_STATE_COMPLETED && agent->last_error) {
        fprintf(stderr, "[Agent] Error: %s\n", agent->last_error);
    }

    agent_destroy(agent);
    return response;
}