// retoor <retoor@molodetz.nl>

#include "context_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define TRUNCATE_THRESHOLD 2000
#define TRUNCATE_KEEP_START 800
#define TRUNCATE_KEEP_END 800
#define TRUNCATE_MARKER "\n\n[... content truncated for context management ...]\n\n"

static bool is_system_message(struct json_object *msg) {
    struct json_object *role_obj;
    if (json_object_object_get_ex(msg, "role", &role_obj)) {
        return strcmp(json_object_get_string(role_obj), "system") == 0;
    }
    return false;
}

static r_status_t truncate_middle(messages_handle msgs, int index) {
    struct json_object *msg = messages_get_object(msgs, index);
    if (!msg) return R_ERROR_NOT_FOUND;

    struct json_object *content_obj;
    const char *content = NULL;
    bool is_tool_result = false;

    if (json_object_object_get_ex(msg, "content", &content_obj)) {
        content = json_object_get_string(content_obj);
    } else if (json_object_object_get_ex(msg, "tool_result", &content_obj)) {
        content = json_object_get_string(content_obj);
        is_tool_result = true;
    }

    if (!content) return R_SUCCESS;

    size_t len = strlen(content);
    if (len <= TRUNCATE_THRESHOLD) return R_SUCCESS;

    char *new_content = malloc(TRUNCATE_KEEP_START + strlen(TRUNCATE_MARKER) + TRUNCATE_KEEP_END + 1);
    if (!new_content) return R_ERROR_OUT_OF_MEMORY;

    strncpy(new_content, content, TRUNCATE_KEEP_START);
    new_content[TRUNCATE_KEEP_START] = '\0';
    strcat(new_content, TRUNCATE_MARKER);
    strcat(new_content, content + len - TRUNCATE_KEEP_END);

    struct json_object *new_msg = json_object_get(msg); // Increments ref count
    if (is_tool_result) {
        json_object_object_add(new_msg, "tool_result", json_object_new_string(new_content));
    } else {
        json_object_object_add(new_msg, "content", json_object_new_string(new_content));
    }

    free(new_content);
    return messages_replace_at(msgs, index, new_msg);
}

r_status_t context_manager_shrink(messages_handle msgs) {
    if (!msgs) return R_ERROR_INVALID_ARG;

    int count = messages_count(msgs);
    if (count <= 1) return R_ERROR_API_ERROR; // Cannot shrink further

    fprintf(stderr, "  \033[2m-> Context limit reached, compressing history...\033[0m\n");

    // Phase 1: Truncate large messages in history (middle-cut)
    // We skip the last message as it's usually the one we just added or the latest prompt
    bool truncated_any = false;
    for (int i = 0; i < count - 1; i++) {
        struct json_object *msg = messages_get_object(msgs, i);
        if (is_system_message(msg)) continue;

        struct json_object *content_obj;
        const char *content = NULL;
        if (json_object_object_get_ex(msg, "content", &content_obj)) {
            content = json_object_get_string(content_obj);
        } else if (json_object_object_get_ex(msg, "tool_result", &content_obj)) {
            content = json_object_get_string(content_obj);
        }

        if (content && strlen(content) > TRUNCATE_THRESHOLD) {
            if (truncate_middle(msgs, i) == R_SUCCESS) {
                truncated_any = true;
            }
        }
    }

    if (truncated_any) return R_SUCCESS;

    // Phase 2: Historical Eviction (remove oldest non-system pairs)
    // We look for the first non-system message
    int first_removable = -1;
    for (int i = 0; i < count - 1; i++) {
        struct json_object *msg = messages_get_object(msgs, i);
        if (!is_system_message(msg)) {
            first_removable = i;
            break;
        }
    }

    if (first_removable != -1 && first_removable < count - 1) {
        // Remove 2 messages to keep user/assistant pairs if possible, 
        // or just one if it's a tool sequence
        int to_remove = (count - first_removable > 2) ? 2 : 1;
        return messages_remove_range(msgs, first_removable, to_remove);
    }

    return R_ERROR_API_ERROR;
}