/* test_agent.c - agent 上下文管理 (不涉及网络) */
#include "agent.h"
#include "json.h"
#include "util.h"

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

static int failures = 0;
static int checks = 0;

#define CHECK(cond, ...)                                                        \
    do {                                                                        \
        checks++;                                                               \
        if (!(cond)) {                                                          \
            failures++;                                                         \
            printf("  ✗ %s:%d: ", __FILE__, __LINE__);                          \
            printf(__VA_ARGS__);                                                \
            printf("\n");                                                       \
        }                                                                       \
    } while (0)

static size_t count_role(const Agent *a, const char *role)
{
    size_t n = 0;
    for (size_t i = 0; i < json_array_size(a->messages); i++) {
        const char *r = json_get_string(json_array_get(a->messages, i), "role");
        if (r && strcmp(r, role) == 0) n++;
    }
    return n;
}

/* ------------------------------------------------------------------ */
/* 上下文超限的自动恢复                                                */
/* ------------------------------------------------------------------ */

static void add_msg(Agent *a, const char *role, const char *content)
{
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", role);
    json_object_set_string(m, "content", content);
    json_array_push(a->messages, m);
}

/* 造一条 assistant 带 tool_calls 的消息 */
static void add_assistant_tool_call(Agent *a, const char *id, const char *name)
{
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", "assistant");
    json_object_set_null(m, "content");
    JsonValue *tcs = json_new_array();
    JsonValue *tc = json_new_object();
    json_object_set_string(tc, "id", id);
    json_object_set_string(tc, "type", "function");
    JsonValue *fn = json_new_object();
    json_object_set_string(fn, "name", name);
    json_object_set_string(fn, "arguments", "{}");
    json_object_set(tc, "function", fn);
    json_array_push(tcs, tc);
    json_object_set(m, "tool_calls", tcs);
    json_array_push(a->messages, m);
}

static void add_tool_result(Agent *a, const char *id, const char *content)
{
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", "tool");
    json_object_set_string(m, "tool_call_id", id);
    json_object_set_string(m, "content", content);
    json_array_push(a->messages, m);
}

/* 序列化历史, 便于断言 */
static char *dump_history(const Agent *a)
{
    return json_dump(a->messages);
}

/* 配对不变量：任何 assistant(tool_calls) 后面必须**紧跟**同 id 的 tool 消息. 
 * 服务端就是这么要求的, 违反了下一轮直接 400 —— 所以每次动过历史之后都要查.  */
static void check_pairing(Agent *a)
{
    size_t n = json_array_size(a->messages);
    for (size_t i = 0; i < n; i++) {
        JsonValue *m = json_array_get(a->messages, i);
        JsonValue *tcs = json_object_get(m, "tool_calls");
        if (!tcs) continue;
        const char *id = json_get_string(json_array_get(tcs, 0), "id");
        CHECK(i + 1 < n, "tool_calls 之后必须有消息");
        if (i + 1 < n) {
            JsonValue *nxt = json_array_get(a->messages, i + 1);
            const char *nid = json_get_string(nxt, "tool_call_id");
            CHECK(nid && id && strcmp(nid, id) == 0,
                  "tool 结果必须紧跟其 tool_call (第 %zu 条)", i);
        }
    }
}

/* 一轮一轮摘下去 (/back 走的就是这条路), 历史在每一步都必须仍然自洽. 
 *
 * 这里原来测的是"上下文超限自动裁剪" (agent_drop_oldest_round)—— 那套机制
 * 已经**整个删掉**：它的判据只是错误文案里的关键词, 却会不可逆地删掉最老的
 * 几轮对话. 现在超限就是如实报错、历史一个字节都不动, 怎么处理由用户决定
 *  (/back 摘一轮、/reset 重来). 但"摘掉一整轮之后 tool 配对不能散"这条
 * 不变量与由谁触发无关, 所以留着继续测.  */
static void test_take_rounds_keeps_pairing(void)
{
    puts("逐轮摘掉之后历史仍然自洽");
    Agent a;
    agent_init(&a);
    agent_add_system(&a, "SYS");
    for (int i = 1; i <= 4; i++) {
        char buf[32];
        snprintf(buf, sizeof buf, "U%d", i);
        add_msg(&a, "user", buf);
        add_assistant_tool_call(&a, "call_x", "read_file");
        add_tool_result(&a, "call_x", "R");
        add_msg(&a, "assistant", "A");
    }

    int rounds = 0;
    for (;;) {
        JsonValue *r = agent_take_last_round(&a);
        if (!r) break;
        json_free(r);
        rounds++;
        check_pairing(&a);        /* 每摘一轮都查一遍配对 */
    }
    CHECK(rounds == 4, "共应摘掉 4 轮, 实际 %d", rounds);
    CHECK(agent_history_len(&a) == 1, "最后应只剩 system, 实际 %zu",
          agent_history_len(&a));
    agent_free(&a);
}

static void test_take_last_round(void)
{
    puts("摘掉最近一轮 (/back 用)");
    Agent a;
    agent_init(&a);
    agent_add_system(&a, "SYS");
    add_msg(&a, "user", "U1");
    add_assistant_tool_call(&a, "call_1", "run_command");
    add_tool_result(&a, "call_1", "R1");
    add_msg(&a, "assistant", "A1");
    add_msg(&a, "user", "U2");
    add_msg(&a, "assistant", "A2");

    CHECK(agent_history_len(&a) == 7, "初始 7 条, 实际 %zu", agent_history_len(&a));

    JsonValue *r = agent_take_last_round(&a);
    CHECK(r != NULL, "应摘到一轮");
    CHECK(r && json_array_size(r) == 2, "最后一轮应 2 条, 实际 %zu",
          r ? json_array_size(r) : 0);
    CHECK(agent_history_len(&a) == 5, "摘完应剩 5 条, 实际 %zu",
          agent_history_len(&a));
    char *h = dump_history(&a);
    CHECK(strstr(h, "U2") == NULL, "U2 应被摘掉");
    CHECK(strstr(h, "U1") != NULL, "U1 应保留");
    free(h);
    json_free(r);

    /* 关键：工具调用与结果必须整组一起走 */
    r = agent_take_last_round(&a);
    CHECK(r && json_array_size(r) == 4, "第一轮应 4 条, 实际 %zu",
          r ? json_array_size(r) : 0);
    h = dump_history(&a);
    CHECK(strstr(h, "call_1") == NULL && strstr(h, "R1") == NULL,
          "工具调用与结果应一并摘掉, 不能只删一半");
    CHECK(strstr(h, "SYS") != NULL, "system 必须保留");
    free(h);
    json_free(r);

    /* 只剩 system 时无可摘 */
    r = agent_take_last_round(&a);
    CHECK(r == NULL, "没有用户轮次时应返回 NULL");
    json_free(r);
    CHECK(agent_history_len(&a) == 1, "历史应只剩 system");
    agent_free(&a);

    /* 没有 user 消息的历史 */
    Agent b;
    agent_init(&b);
    add_msg(&b, "assistant", "只有助手");
    CHECK(agent_take_last_round(&b) == NULL, "没有 user 时不应摘任何东西");
    CHECK(agent_history_len(&b) == 1, "历史不应被改动");
    agent_free(&b);
}

/* ------------------------------------------------------------------ */
/* 结构体检：assistant(tool_calls) 与 tool 消息的配对                   */
/* ------------------------------------------------------------------ */

/* 造一条 assistant, 带 count 个工具调用 (id 依次是 c0, c1, ...) */
static JsonValue *mk_assistant(size_t count)
{
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", "assistant");
    JsonValue *tcs = json_new_array();
    for (size_t i = 0; i < count; i++) {
        char id[16];
        snprintf(id, sizeof id, "c%zu", i);
        JsonValue *tc = json_new_object();
        json_object_set_string(tc, "id", id);
        json_object_set_string(tc, "type", "function");
        JsonValue *fn = json_new_object();
        json_object_set_string(fn, "name", "run_command");
        json_object_set_string(fn, "arguments", "{}");
        json_object_set(tc, "function", fn);
        json_array_push(tcs, tc);
    }
    json_object_set(m, "tool_calls", tcs);
    return m;
}

static JsonValue *mk_tool(const char *id)
{
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", "tool");
    json_object_set_string(m, "tool_call_id", id);
    json_object_set_string(m, "content", "ok");
    return m;
}

static JsonValue *mk_plain(const char *role, const char *text)
{
    JsonValue *m = json_new_object();
    json_object_set_string(m, "role", role);
    json_object_set_string(m, "content", text);
    return m;
}

/* 数一数某个 assistant 之后紧跟的 tool 消息覆盖了几个 id */
static const char *nth_role(const JsonValue *msgs, size_t i)
{
    return json_get_string(json_array_get((JsonValue *)msgs, i), "role");
}

static int pairing_ok(const JsonValue *msgs)
{
    size_t n = json_array_size((JsonValue *)msgs);
    for (size_t i = 0; i < n; i++) {
        JsonValue *m = json_array_get((JsonValue *)msgs, i);
        const char *r = json_get_string(m, "role");
        if (!r || strcmp(r, "assistant") != 0) continue;
        JsonValue *tcs = json_object_get(m, "tool_calls");
        if (!tcs || json_array_size(tcs) == 0) continue;
        for (size_t k = 0; k < json_array_size(tcs); k++) {
            const char *id = json_get_string(json_array_get(tcs, k), "id");
            size_t j = i + 1;
            int found = 0;
            for (; j < n && strcmp(nth_role(msgs, j), "tool") == 0; j++)
                if (strcmp(json_get_string(json_array_get((JsonValue *)msgs, j),
                                           "tool_call_id"), id) == 0) found = 1;
            if (!found) return 0;
        }
    }
    return 1;
}

static void test_repair_pairing(void)
{
    int fixed = 0;

    puts("结构体检：本来没问题就不该动它");
    JsonValue *ok = json_new_array();
    json_array_push(ok, mk_plain("user", "跑一下"));
    json_array_push(ok, mk_assistant(2));
    json_array_push(ok, mk_tool("c0"));
    json_array_push(ok, mk_tool("c1"));
    json_array_push(ok, mk_plain("assistant", "做完了"));
    CHECK(pairing_ok(ok), "构造出来的历史本身应是合法的");
    CHECK(agent_repair_tool_pairing(ok, &fixed) == NULL,
          "合法历史不该被重建 (NULL = 什么都没改)");
    CHECK(fixed == 0, "合法历史的修复处数应为 0");
    json_free(ok);

    puts("缺 tool 消息：补齐, 且补在原位");
    /* 复现真实报错：assistant 带两个调用, 中间夹了一条 user (图省事直接
     * 模仿当年那个 bug 的形状), 第二个调用没人回应 */
    JsonValue *bad = json_new_array();
    json_array_push(bad, mk_plain("user", "看看图"));
    json_array_push(bad, mk_assistant(2));
    json_array_push(bad, mk_tool("c0"));
    json_array_push(bad, mk_plain("user", "图片附件"));
    json_array_push(bad, mk_tool("c1"));
    fixed = 0;
    JsonValue *r = agent_repair_tool_pairing(bad, &fixed);
    CHECK(r != NULL, "对不上时应该返回重建好的数组");
    if (r) {
        CHECK(pairing_ok(r), "修完之后配对必须成立");
        /* 两处：给 c1 补一条占位 (它没被回应), 再把错位的那条 c1 结果丢掉
         *  (它被 user 消息隔开, 已经不在 tool 段里了 —— 留着就是孤儿). 
         * 丢掉真实结果听着可惜, 但结构坏了只能二选一：要么丢它, 要么留一段
         * 服务端必然拒绝的历史. 模型会从占位消息里知道该重跑.  */
        CHECK(fixed == 2, "应补 1 处、丢 1 处孤儿, 实际 %d", fixed);
        CHECK(json_array_size(r) == 5,
              "补一条、丢一条之后仍是 5 条, 实际 %zu", json_array_size(r));
        /* 占位消息必须紧跟在 assistant 的 tool 段里, 而不是被塞到最后 */
        CHECK(strcmp(nth_role(r, 2), "tool") == 0 &&
              strcmp(nth_role(r, 3), "tool") == 0,
              "两条 tool 消息应紧跟在 assistant 之后");
        const char *c = json_get_string(json_array_get(r, 3), "content");
        CHECK(c && strstr(c, "没有执行") == NULL && strstr(c, "没有得到结果") != NULL,
              "占位消息要说明它没拿到结果, 实际 %s", c ? c : "(null)");
        json_free(r);
    }
    json_free(bad);

    puts("末尾悬空：assistant 之后再没有消息");
    JsonValue *tail = json_new_array();
    json_array_push(tail, mk_plain("user", "跑一下"));
    json_array_push(tail, mk_assistant(1));
    fixed = 0;
    r = agent_repair_tool_pairing(tail, &fixed);
    CHECK(r && pairing_ok(r) && fixed == 1, "结尾缺的也要补上 (fixed=%d)", fixed);
    json_free(r);
    json_free(tail);

    puts("孤儿 tool 消息：没有对应的 tool_calls, 要丢掉");
    JsonValue *orphan = json_new_array();
    json_array_push(orphan, mk_plain("user", "你好"));
    json_array_push(orphan, mk_tool("谁也不是"));
    JsonValue *r2 = agent_repair_tool_pairing(orphan, &fixed);
    CHECK(fixed == 1, "孤儿 tool 消息应被计为 1 处, 实际 %d", fixed);
    CHECK(r2 != NULL && json_array_size(r2) == 1, "孤儿应被丢掉, 只剩 user");
    json_free(r2);
    json_free(orphan);

    puts("空输入与边界");
    fixed = 7;
    CHECK(agent_repair_tool_pairing(NULL, &fixed) == NULL && fixed == 0, "NULL 输入");
    JsonValue *empty = json_new_array();
    CHECK(agent_repair_tool_pairing(empty, &fixed) == NULL, "空历史没什么可修");
    json_free(empty);
    /* out_fixed 可以为 NULL */
    JsonValue *b2 = json_new_array();
    json_array_push(b2, mk_assistant(1));
    JsonValue *r3 = agent_repair_tool_pairing(b2, NULL);
    json_free(r3);
    json_free(b2);
    CHECK(1, "out_fixed 为 NULL 时不崩");
}

/* ------------------------------------------------------------------ */
/* 面向用户的历史口径：不含 system                                      */
/* ------------------------------------------------------------------ */

static void test_conversation_len(void)
{
    puts("历史条数：system 是配置, 不该算进对话");
    Agent a;
    agent_init(&a);

    CHECK(agent_history_len(&a) == 0, "空 agent 的数组长度是 0");
    CHECK(agent_conversation_len(&a) == 0, "空 agent 的对话长度也是 0");

    /* 只有 system：数组里 1 条, 但对话是空的 */
    agent_add_system(&a, "你是助手");
    CHECK(agent_history_len(&a) == 1, "只有 system 时数组长度 1");
    CHECK(agent_conversation_len(&a) == 0,
          "只有 system 时**对话**长度必须是 0 (用户报的就是这里)");

    add_msg(&a, "user", "hi");
    add_msg(&a, "assistant", "你好");
    CHECK(agent_history_len(&a) == 3, "system + 一问一答 = 3");
    CHECK(agent_conversation_len(&a) == 2, "对话长度 2");

    /* 没有 system 的历史 (比如裸构造的)不该被误减 */
    Agent b;
    agent_init(&b);
    add_msg(&b, "user", "hi");
    CHECK(agent_history_len(&b) == 1 && agent_conversation_len(&b) == 1,
          "没有 system 时两个口径一致");

    /* reset 之后回到"只有 system" */
    agent_reset(&a);
    CHECK(agent_history_len(&a) == 1 && agent_conversation_len(&a) == 0,
          "reset 之后数组剩 system、对话为空");

    agent_free(&a);
    agent_free(&b);
}

int main(void)
{
    Agent a;
    agent_init(&a);
    CHECK(a.messages && a.messages->type == JSON_ARRAY, "messages 应为数组");

    agent_add_system(&a, "S1");
    CHECK(agent_history_len(&a) == 1, "应只有 1 条 system");
    CHECK(count_role(&a, "system") == 1, "system 计数");

    agent_add_user(&a, "U1");
    CHECK(agent_history_len(&a) == 2, "应 2 条");

    /* 关键不变量：重复设置 system 应替换而不是追加 */
    agent_add_system(&a, "S2");
    CHECK(count_role(&a, "system") == 1, "重复设置 system 不应叠加, 实际 %zu 条",
          count_role(&a, "system"));
    CHECK(strcmp(json_get_string(json_array_get(a.messages, 0), "content"), "S2") == 0,
          "新 system 应在第 0 条");
    CHECK(strcmp(json_get_string(json_array_get(a.messages, 1), "content"), "U1") == 0,
          "原有对话应保留");

    agent_add_user(&a, "U2");
    CHECK(agent_history_len(&a) == 3, "应 3 条");

    agent_reset(&a);
    CHECK(agent_history_len(&a) == 1, "reset 后只保留 system, 实际 %zu",
          agent_history_len(&a));
    CHECK(count_role(&a, "system") == 1, "reset 后 system 仍在");
    CHECK(strcmp(json_get_string(json_array_get(a.messages, 0), "content"), "S2") == 0,
          "reset 保留的是最新 system");

    /* 历史里没有 system 时 reset 应得到空历史 */
    Agent b;
    agent_init(&b);
    agent_add_user(&b, "only user");
    agent_reset(&b);
    CHECK(agent_history_len(&b) == 0, "无 system 时 reset 应为空, 实际 %zu",
          agent_history_len(&b));

    agent_free(&a);
    agent_free(&b);
    test_take_rounds_keeps_pairing();
    test_take_last_round();
    test_repair_pairing();
    test_conversation_len();
    printf("\n%d 项检查, %d 项失败\n", checks, failures);
    return failures ? 1 : 0;
}