#include "compact.h"
#include "message.h"

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

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

#define CHECK(cond, msg)                                                     \
    do {                                                                     \
        checks++;                                                            \
        if (!(cond)) {                                                       \
            failures++;                                                      \
            fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, (msg)); \
        }                                                                    \
    } while (0)

#define CHECK_EQ_U64(got, want, msg)                                         \
    do {                                                                     \
        checks++;                                                            \
        if ((unsigned long long)(got) != (unsigned long long)(want)) {       \
            failures++;                                                      \
            fprintf(stderr, "FAIL %s:%d: %s (got %llu, want %llu)\n",       \
                    __FILE__, __LINE__, (msg), (unsigned long long)(got),   \
                    (unsigned long long)(want));                             \
        }                                                                    \
    } while (0)

#define CHECK_CONT(hay, needle, msg)                                         \
    do {                                                                     \
        checks++;                                                            \
        if (!hay || !strstr(hay, needle)) {                                  \
            failures++;                                                      \
            fprintf(stderr, "FAIL %s:%d: %s (got `%s`, want substring `%s`)\n", \
                    __FILE__, __LINE__, (msg), hay ? hay : "(null)",         \
                    (needle));                                               \
        }                                                                    \
    } while (0)

static char *repeat_char(char c, size_t n) {
    char *s = (char* ) xmalloc(n + 1);
    memset(s, c, n);
    s[n] = '\0';
    return s;
}

/* text of the first tool-result item of a user message (or ""). */
static char *trc_text_of(const message *m) {
    if (!m->is_user() || m->user().parts.empty())
        return xstrdup("");
    const user_part &p = m->user().parts[0];
    const user_tool_result_part *r = std::get_if<user_tool_result_part>(&p);
    if (!r || r->content.empty())
        return xstrdup("");
    const trc_text *t = std::get_if<trc_text>(&r->content[0]);
    return xstrdup(t ? t->text.c_str() : "");
}

/* ------------------------------------------------------------------ */
/* estimate_tokens                                                     */
/* ------------------------------------------------------------------ */

static void test_estimate_tokens(void) {
    char *body = repeat_char('a', 40);
    message m = msg_user_text(body);
    CHECK_EQ_U64(compact_estimate_tokens(&m, 1), 10,
                 "40 bytes -> 10 tokens");
    free(body);

    CHECK_EQ_U64(compact_estimate_tokens(NULL, 0), 0, "empty -> 0 tokens");
}

static void test_estimate_tokens_tool_call_args(void) {
    message m = msg_assistant_tool_call("read_file", "{\"path\":\"a\"}");
    CHECK_EQ_U64(compact_estimate_tokens(&m, 1), 6,
                 "name(9)+args(12)=21 -> 6 tokens");
}

/* ------------------------------------------------------------------ */
/* keep_boundary                                                       */
/* ------------------------------------------------------------------ */

static void test_keep_boundary(void) {
    message msgs[6] = {
        msg_user_text("read the file"),
        msg_assistant_text("ok"),
        msg_user_tool_result("c1", "read_file", "contents"),
        msg_assistant_text("here it is"),
        msg_user_text("now edit it"),
        msg_assistant_text("done"),
    };
    CHECK_EQ_U64(compact_keep_boundary(msgs, 6), 4,
                 "boundary is last user");

    message n[2] = {
        msg_user_tool_result("c1", "read_file", "contents"),
        msg_assistant_text("done"),
    };
    CHECK_EQ_U64(compact_keep_boundary(n, 2), 2, "none -> keep all");

    CHECK_EQ_U64(compact_keep_boundary(NULL, 0), 0, "empty -> 0");
}

/* ------------------------------------------------------------------ */
/* clear_old_tool_results                                              */
/* ------------------------------------------------------------------ */

static void test_clear_old_tool_results(void) {
    message msgs[6] = {
        msg_user_text("read the file"),
        msg_assistant_text("ok"),
        msg_user_tool_result("c1", "read_file", "lots of content"),
        msg_assistant_text("here it is"),
        msg_user_text("now edit it"),
        msg_user_tool_result("c2", "write_file", "tail content"),
    };
    size_t boundary = compact_keep_boundary(msgs, 6);
    CHECK_EQ_U64(boundary, 4, "boundary == 4");

    size_t cleared = compact_clear_old_tool_results(msgs, 6, boundary);
    CHECK_EQ_U64(cleared, 1, "cleared exactly one");

    char *t2 = trc_text_of(&msgs[2]);
    CHECK(xstreq(t2, CLEARED_PLACEHOLDER), "stale result placeholdered");
    free(t2);

    char *t5 = trc_text_of(&msgs[5]);
    CHECK(xstreq(t5, "tail content"), "tail result untouched");
    free(t5);

    /* name + call preserved on the cleared result */
    CHECK(msgs[2].is_user(), "still a user message");
    CHECK(std::holds_alternative<user_tool_result_part>(msgs[2].user().parts[0]),
          "still a tool result");
    CHECK(std::get<user_tool_result_part>(msgs[2].user().parts[0]).name ==
              "read_file",
          "name preserved");
    CHECK(std::get<user_tool_result_part>(msgs[2].user().parts[0]).call ==
              "c1",
          "call id preserved");

    /* idempotent */
    CHECK_EQ_U64(compact_clear_old_tool_results(msgs, 6, boundary), 0,
                 "second pass clears nothing");
}

/* ------------------------------------------------------------------ */
/* transcript_markdown                                                 */
/* ------------------------------------------------------------------ */

static void test_transcript_markdown(void) {
    message msgs[4] = {
        msg_user_text("read the file"),
        msg_assistant_tool_call("read_file", "{\"path\":\"a.txt\"}"),
        msg_user_tool_result("c1", "read_file", "line one\nline two"),
        msg_assistant_text("The file has two lines."),
    };
    char *out = compact_transcript_markdown(msgs, 4);
    CHECK(out && strncmp(out, "# llsh transcript", 15) == 0, "header");
    CHECK_CONT(out, "## User\n\nread the file", "user block");
    CHECK_CONT(out, "### Tool call: `read_file`", "tool call heading");
    CHECK_CONT(out, "```json", "json fence");
    CHECK_CONT(out, "\"path\"", "json path key");
    CHECK_CONT(out, "### Tool result: `read_file`", "tool result heading");
    CHECK_CONT(out, "line one\nline two", "result body");
    CHECK_CONT(out, "## Assistant\n\nThe file has two lines.", "assistant");
    free(out);
}

static void test_transcript_markdown_empty(void) {
    char *out = compact_transcript_markdown(NULL, 0);
    CHECK(out && strncmp(out, "# llsh transcript", 15) == 0, "header");
    CHECK_CONT(out, "_No messages yet._", "placeholder");
    free(out);
}

/* ------------------------------------------------------------------ */
/* recent_file_paths                                                   */
/* ------------------------------------------------------------------ */

static void test_recent_file_paths(void) {
    strlist out;
    const char *s1 = "## Goal\ndo a thing\n\n## Files\n- src/foo.rs\n"
                     "* src/bar.rs\n\n## Next steps\n- more";
    compact_recent_file_paths(s1, &out);
    CHECK(out.len == 2, "two paths");
    if (out.len == 2) {
        CHECK(xstreq(out.items[0], "src/foo.rs"), "path 0");
        CHECK(xstreq(out.items[1], "src/bar.rs"), "path 1");
    }
    sl_free(&out);

    compact_recent_file_paths("## Goal\nnothing here", &out);
    CHECK(out.len == 0, "no files section -> empty");
    sl_free(&out);

    const char *s3 = "## Next steps\n- src/foo.rs\n\n## Files\n- src/bar.rs";
    compact_recent_file_paths(s3, &out);
    CHECK(out.len == 1, "ignores bullets outside files");
    if (out.len == 1)
        CHECK(xstreq(out.items[0], "src/bar.rs"), "only bar.rs");
    sl_free(&out);
}

/* ------------------------------------------------------------------ */
/* compose_history                                                     */
/* ------------------------------------------------------------------ */

static void test_compose_history(void) {
    message tail[2] = {
        msg_user_text("now edit it"),
        msg_assistant_text("done"),
    };
    const char *paths[1] = {"src/foo.rs"};
    const char *contents[1] = {"fn main() {}"};
    history *h = compact_compose_history("## Goal\nship it", paths, contents,
                                         1, tail, 2);
    CHECK_EQ_U64(history_len(h), 3, "one body + tail");
    const message *first = history_at(h, 0);
    CHECK(first->is_user(), "first is user");
    CHECK_EQ_U64(first->user().parts.size(), 1, "single text part");
    const std::string &body =
        std::get<user_text_part>(first->user().parts[0]).text;
    CHECK_CONT(body.c_str(), "## Goal", "summary present");
    CHECK_CONT(body.c_str(), "src/foo.rs", "file path present");
    CHECK_CONT(body.c_str(), "fn main() {}", "file content present");
    CHECK(history_at(h, 1)->is_user(), "tail[0] user");
    CHECK(history_at(h, 2)->is_assistant(), "tail[1] assistant");
    history_free(h);
    free(h);
}

static void test_compose_history_caps(void) {
    char *big = repeat_char('x', 200000);
    const char *paths[1] = {"big.log"};
    const char *contents[1] = {big};
    history *h =
        compact_compose_history("summary", paths, contents, 1, NULL, 0);
    CHECK_EQ_U64(history_len(h), 1, "just the body");
    const std::string &body =
        std::get<user_text_part>(history_at(h, 0)->user().parts[0]).text;
    CHECK(utf8_len(body.c_str()) < 200000, "content capped, not full");
    char *want = repeat_char('x', PER_FILE_CAP);
    CHECK_CONT(body.c_str(), want, "first PER_FILE_CAP chars present");
    free(want);
    history_free(h);
    free(h);
    free(big);
}

/* ------------------------------------------------------------------ */
/* transcript (flat)                                                   */
/* ------------------------------------------------------------------ */

static void test_transcript_flat(void) {
    message msgs[3] = {
        msg_user_text("read the file"),
        msg_assistant_text("reading"),
        msg_user_tool_result("c1", "read_file", "the body"),
    };
    char *out = compact_transcript(msgs, 3);
    CHECK_CONT(out, "User: read the file", "user line");
    CHECK_CONT(out, "Assistant: reading", "assistant line");
    CHECK_CONT(out, "ToolResult(read_file): the body", "tool result line");
    free(out);
}

/* ------------------------------------------------------------------ */
/* compact_at / clear_at                                               */
/* ------------------------------------------------------------------ */

static void test_thresholds(void) {
    compact_config c;
    compact_config_from(1000, 0.75, 0.5, &c);
    unsigned long v;
    CHECK(compact_at(&c, &v), "compact enabled");
    CHECK_EQ_U64(v, 750, "compact_at == 750");
    CHECK(clear_at(&c, &v), "clear enabled");
    CHECK_EQ_U64(v, 500, "clear_at == 500");

    compact_config d;
    compact_config_from(1000, 0.0, 0.0, &d);
    CHECK(!compact_at(&d, &v), "compact disabled");
    CHECK(!clear_at(&d, &v), "clear disabled");
}

/* ------------------------------------------------------------------ */
/* message model                                                       */
/* ------------------------------------------------------------------ */

static void test_message_model(void) {
    message m = msg_user_tool_result("c9", "grep", "match");
    CHECK(m.is_user(), "role user");
    CHECK_EQ_U64(m.user().parts.size(), 1, "one part");
    CHECK(std::holds_alternative<user_tool_result_part>(m.user().parts[0]),
          "tool result kind");
    CHECK(std::get<user_tool_result_part>(m.user().parts[0]).name == "grep",
          "name");
    CHECK(std::get<user_tool_result_part>(m.user().parts[0]).call == "c9",
          "call");
    CHECK(std::get<trc_text>(
              std::get<user_tool_result_part>(m.user().parts[0]).content[0])
              .text == "match",
          "content");

    /* copy round-trip */
    message orig = msg_assistant_tool_call("write_file", "{\"path\":\"a\"}");
    message copy = orig;
    CHECK(copy.is_assistant(), "copy role");
    CHECK_EQ_U64(copy.assistant().parts.size(), 1, "copy len");
    CHECK(std::get<asst_tool_call_part>(copy.assistant().parts[0]).name ==
              "write_file",
          "copy name");
    CHECK(std::get<asst_tool_call_part>(copy.assistant().parts[0]).args ==
              "{\"path\":\"a\"}",
          "copy args");
    /* independent storage: mutating orig does not affect copy */
    std::get<asst_tool_call_part>(orig.assistant().parts[0]).name = "other";
    CHECK(std::get<asst_tool_call_part>(copy.assistant().parts[0]).name ==
              "write_file",
          "copy independent");

    /* multi-part assistant message */
    message multi = msg_assistant_text("thinking...");
    msg_push_assistant_tool_call(multi, "run_shell", "echo hi");
    msg_push_assistant_reasoning(multi, "because");
    CHECK_EQ_U64(multi.assistant().parts.size(), 3, "three parts");
    CHECK(std::holds_alternative<asst_text_part>(multi.assistant().parts[0]),
          "part0 text");
    CHECK(std::holds_alternative<asst_tool_call_part>(
              multi.assistant().parts[1]),
          "part1 tool call");
    CHECK(std::holds_alternative<asst_reasoning_part>(
              multi.assistant().parts[2]),
          "part2 reasoning");

    /* history push / take */
    history h;
    history_push(&h, msg_user_text("one"));
    history_push(&h, msg_user_text("two"));
    history_push(&h, msg_user_text("three"));
    CHECK_EQ_U64(history_len(&h), 3, "len 3");
    message mid = history_take(&h, 1);
    CHECK(std::get<user_text_part>(mid.user().parts[0]).text == "two",
          "take returns elem 1");
    CHECK_EQ_U64(history_len(&h), 2, "len 2 after take");
    CHECK(std::get<user_text_part>(history_at(&h, 0)->user().parts[0]).text ==
              "one",
          "elem0 intact");
    CHECK(std::get<user_text_part>(history_at(&h, 1)->user().parts[0]).text ==
              "three",
          "elem1 shifted");
    /* stack history: RAII destructor cleans up at scope exit */

    /* preamble sanity */
    const char *p = compact_summarizer_preamble();
    CHECK_CONT(p, "## Goal", "preamble goal section");
    CHECK_CONT(p, "## Files", "preamble files section");
    CHECK_CONT(p, "under 800 words", "preamble word limit");
}

/* ------------------------------------------------------------------ */
/* rubric (SelfCompact)                                                */
/* ------------------------------------------------------------------ */

static void test_rubric_parse_valid(void) {
    const char *answer =
        "C1: Y -- the last message ends with a finished tool result\n"
        "C2: Y -- 1. fact one \"x\"\n"
        "C3: Y -- a new URL was obtained\n"
        "N1: N -- each recent call returned new content\n";
    compact_rubric r;
    char *err = NULL;
    CHECK(compact_rubric_parse(answer, &r, &err) == 0, "parse ok");
    CHECK(err == NULL, "no err");
    CHECK_EQ_U64(r.c1, 1, "c1 Y");
    CHECK_EQ_U64(r.c2, 1, "c2 Y");
    CHECK_EQ_U64(r.c3, 1, "c3 Y");
    CHECK_EQ_U64(r.n1, 0, "n1 N");
    CHECK(compact_rubric_fire(&r), "fires when c1&c2&c3 & !n1");
}

static void test_rubric_parse_n(void) {
    const char *answer =
        "C1: N -- still mid-thought\n"
        "C2: Y -- facts\n"
        "C3: Y -- progress\n"
        "N1: N -- fresh results\n";
    compact_rubric r;
    char *err = NULL;
    CHECK(compact_rubric_parse(answer, &r, &err) == 0, "parse ok");
    CHECK_EQ_U64(r.c1, 0, "c1 N");
    CHECK(!compact_rubric_fire(&r), "does not fire with c1 N");
}

static void test_rubric_parse_stuck(void) {
    const char *answer =
        "C1: Y -- closed\n"
        "C2: Y -- facts\n"
        "C3: Y -- progress\n"
        "N1: Y -- three of last four calls repeated\n";
    compact_rubric r;
    char *err = NULL;
    CHECK(compact_rubric_parse(answer, &r, &err) == 0, "parse ok");
    CHECK_EQ_U64(r.n1, 1, "n1 Y");
    CHECK(!compact_rubric_fire(&r), "does not fire when stuck");
}

static void test_rubric_parse_missing(void) {
    const char *answer =
        "C1: Y -- closed\n"
        "C2: Y -- facts\n"
        "C3: Y -- progress\n";
    compact_rubric r;
    char *err = NULL;
    CHECK(compact_rubric_parse(answer, &r, &err) != 0, "missing N1 -> error");
    CHECK(err != NULL, "err set");
    free(err);
}

static void test_rubric_parse_malformed(void) {
    const char *answer =
        "C1: MAYBE -- closed\n"
        "C2: Y -- facts\n"
        "C3: Y -- progress\n"
        "N1: N -- fresh\n";
    compact_rubric r;
    char *err = NULL;
    CHECK(compact_rubric_parse(answer, &r, &err) != 0, "bad token -> error");
    CHECK(err != NULL, "err set");
    free(err);
}

static void test_rubric_y_without_evidence_is_n(void) {
    /* A Y with no evidence on the line is downgraded to N (the paper's rule). */
    const char *answer =
        "C1: Y\n"
        "C2: Y -- facts\n"
        "C3: Y -- progress\n"
        "N1: N -- fresh\n";
    compact_rubric r;
    char *err = NULL;
    CHECK(compact_rubric_parse(answer, &r, &err) == 0, "parse ok");
    CHECK_EQ_U64(r.c1, 0, "Y without evidence downgraded to N");
    CHECK(!compact_rubric_fire(&r), "does not fire");
}

static void test_rubric_prompt_and_preamble(void) {
    const char *p = compact_rubric_prompt();
    CHECK_CONT(p, "C1", "prompt has C1");
    CHECK_CONT(p, "C2", "prompt has C2");
    CHECK_CONT(p, "C3", "prompt has C3");
    CHECK_CONT(p, "N1", "prompt has N1");
    CHECK_CONT(p, "C1: Y/N -- <evidence>", "prompt output format");

    const char *pre = compact_rubric_preamble();
    CHECK_CONT(pre, "judge", "preamble is a judge role");
}

static void test_rubric_due(void) {
    /* all gates satisfied */
    CHECK(compact_rubric_due(RUBRIC_MIN_ITERATION, RUBRIC_MIN_TOKENS, 0,
                             RUBRIC_DEFAULT_MAX_SUMMARIES,
                             RUBRIC_MIN_ROUNDS_SINCE_PROBE),
          "due at exact thresholds");
    /* iteration too low */
    CHECK(!compact_rubric_due(RUBRIC_MIN_ITERATION - 1, RUBRIC_MIN_TOKENS, 0,
                              RUBRIC_DEFAULT_MAX_SUMMARIES,
                              RUBRIC_MIN_ROUNDS_SINCE_PROBE),
          "not due: iteration");
    /* tokens too low */
    CHECK(!compact_rubric_due(RUBRIC_MIN_ITERATION, RUBRIC_MIN_TOKENS - 1, 0,
                              RUBRIC_DEFAULT_MAX_SUMMARIES,
                              RUBRIC_MIN_ROUNDS_SINCE_PROBE),
          "not due: tokens");
    /* cap reached */
    CHECK(!compact_rubric_due(RUBRIC_MIN_ITERATION, RUBRIC_MIN_TOKENS, 1,
                              RUBRIC_DEFAULT_MAX_SUMMARIES,
                              RUBRIC_MIN_ROUNDS_SINCE_PROBE),
          "not due: cap");
    /* rounds since probe too low */
    CHECK(!compact_rubric_due(RUBRIC_MIN_ITERATION, RUBRIC_MIN_TOKENS, 0,
                              RUBRIC_DEFAULT_MAX_SUMMARIES,
                              RUBRIC_MIN_ROUNDS_SINCE_PROBE - 1),
          "not due: rounds");
    /* unlimited cap (0) */
    CHECK(compact_rubric_due(RUBRIC_MIN_ITERATION, RUBRIC_MIN_TOKENS, 5, 0,
                             RUBRIC_MIN_ROUNDS_SINCE_PROBE),
          "due: unlimited cap");
}

int main(void) {
    test_estimate_tokens();
    test_estimate_tokens_tool_call_args();
    test_keep_boundary();
    test_clear_old_tool_results();
    test_transcript_markdown();
    test_transcript_markdown_empty();
    test_recent_file_paths();
    test_compose_history();
    test_compose_history_caps();
    test_transcript_flat();
    test_thresholds();
    test_message_model();
    test_rubric_parse_valid();
    test_rubric_parse_n();
    test_rubric_parse_stuck();
    test_rubric_parse_missing();
    test_rubric_parse_malformed();
    test_rubric_y_without_evidence_is_n();
    test_rubric_prompt_and_preamble();
    test_rubric_due();
    printf("%d checks, %d failures\n", checks, failures);
    return failures ? 1 : 0;
}