// retoor <retoor@molodetz.nl>

#ifndef R_AGENT_H
#define R_AGENT_H

#include "chat.h"
#include "http_curl.h"
#include "messages.h"
#include "r.h"
#include "tools.h"
#include <stdbool.h>
#include <string.h>
#include <time.h>

#define AGENT_MAX_ITERATIONS 300
#define AGENT_MAX_TOOL_RETRIES 3

typedef enum {
  AGENT_STATUS_RUNNING,
  AGENT_STATUS_COMPLETED,
  AGENT_STATUS_MAX_ITERATIONS,
  AGENT_STATUS_ERROR,
  AGENT_STATUS_TOOL_ERROR
} agent_status_t;

typedef struct {
  char *goal;
  int iteration_count;
  int max_iterations;
  int tool_retry_count;
  int max_tool_retries;
  agent_status_t status;
  time_t start_time;
  char *last_error;
  bool verbose;
} agent_state_t;

static agent_state_t *current_agent = NULL;

agent_state_t *agent_create(const char *goal) {
  agent_state_t *agent = (agent_state_t *)malloc(sizeof(agent_state_t));
  if (agent == NULL) {
    return NULL;
  }

  agent->goal = goal ? strdup(goal) : NULL;
  agent->iteration_count = 0;
  agent->max_iterations = AGENT_MAX_ITERATIONS;
  agent->tool_retry_count = 0;
  agent->max_tool_retries = AGENT_MAX_TOOL_RETRIES;
  agent->status = AGENT_STATUS_RUNNING;
  agent->start_time = time(NULL);
  agent->last_error = NULL;
  agent->verbose = is_verbose;

  return agent;
}

void agent_destroy(agent_state_t *agent) {
  if (agent == NULL) {
    return;
  }
  if (agent->goal) {
    free(agent->goal);
  }
  if (agent->last_error) {
    free(agent->last_error);
  }
  free(agent);
}

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

static struct json_object *agent_process_response(const char *api_url,
                                                   const char *json_data) {
  char *response = curl_post(api_url, json_data);
  if (!response) {
    return NULL;
  }

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

  if (!parsed_json) {
    return NULL;
  }

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

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

  struct json_object *first_choice = json_object_array_get_idx(choices_array, 0);
  if (!first_choice) {
    json_object_put(parsed_json);
    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 const char *agent_get_finish_reason(struct json_object *choice) {
  struct json_object *finish_reason_obj;
  if (json_object_object_get_ex(choice, "finish_reason", &finish_reason_obj)) {
    return json_object_get_string(finish_reason_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 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 bool agent_response_indicates_incomplete(const char *content) {
  if (content == NULL) {
    return false;
  }

  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
  };

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

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

  size_t len = strlen(content);
  if (len > 3) {
    for (int i = 0; incomplete_endings[i] != NULL; 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_state_t *agent, const char *user_message) {
  if (agent == NULL) {
    return NULL;
  }

  current_agent = agent;
  agent->status = AGENT_STATUS_RUNNING;
  agent->iteration_count = 0;
  agent->tool_retry_count = 0;

  if (user_message == NULL || *user_message == '\0') {
    agent->status = AGENT_STATUS_ERROR;
    agent_set_error(agent, "Empty user message");
    return NULL;
  }

  char *json_data = chat_json("user", user_message);
  if (json_data == NULL) {
    agent->status = AGENT_STATUS_ERROR;
    agent_set_error(agent, "Failed to create chat JSON");
    return NULL;
  }

  char *final_response = NULL;

  while (agent->status == AGENT_STATUS_RUNNING) {
    agent->iteration_count++;

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

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

    struct json_object *choice =
        agent_process_response(get_completions_api_url(), json_data);

    if (choice == NULL) {
      agent->tool_retry_count++;
      if (agent->tool_retry_count >= agent->max_tool_retries) {
        agent->status = AGENT_STATUS_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);
      }
      continue;
    }

    agent->tool_retry_count = 0;

    struct json_object *message_obj = agent_get_message(choice);
    if (message_obj) {
      message_add_object(json_object_get(message_obj));
    }

    const char *finish_reason = agent_get_finish_reason(choice);
    bool has_tools = agent_has_tool_calls(choice);

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

    if (has_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 *tool_results = tools_execute(tool_calls);

      int results_count = json_object_array_length(tool_results);
      for (int i = 0; i < results_count; i++) {
        struct json_object *tool_result =
            json_object_array_get_idx(tool_results, i);
        message_add_tool_call(json_object_get(tool_result));
      }

      json_data = chat_json(NULL, NULL);
      if (json_data == NULL) {
        agent->status = AGENT_STATUS_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 = chat_json("user", "Continue. Execute the necessary actions to complete the task.");
        if (json_data == NULL) {
          agent->status = AGENT_STATUS_ERROR;
          agent_set_error(agent, "Failed to create continue JSON");
          break;
        }

      } else {
        final_response = content;
        agent->status = AGENT_STATUS_COMPLETED;

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

  current_agent = NULL;
  return final_response;
}

char *agent_chat(const char *user_message) {
  agent_state_t *agent = agent_create(user_message);
  if (agent == NULL) {
    return NULL;
  }

  char *response = agent_run(agent, user_message);

  if (agent->verbose && agent->status != AGENT_STATUS_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) {
  agent_state_t *agent = agent_create(user_message);
  if (agent == NULL) {
    return NULL;
  }

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

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

  agent_destroy(agent);
  return response;
}

#endif