#include "compact.h"

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

/* ------------------------------------------------------------------ */
/* Preamble                                                            */
/* ------------------------------------------------------------------ */

const char *compact_summarizer_preamble(void) {
    return "You are a context compactor for a terminal coding agent. You are "
           "given the transcript of the earlier part of that agent's "
           "conversation, which no longer fits in the context window. Compress "
           "it into a faithful, dense summary that lets the agent continue "
           "seamlessly.\n"
           "\n"
           "Write ONLY the summary, in Markdown, with these exact sections in "
           "this order:\n"
           "\n"
           "## Goal\n"
           "What the user is trying to accomplish.\n"
           "\n"
           "## Decisions\n"
           "Choices made and the reasons that still matter.\n"
           "\n"
           "## Work done\n"
           "Concrete steps taken: files touched, commands run, and their "
           "results.\n"
           "\n"
           "## Open issues\n"
           "Errors, edge cases, or problems that remain unresolved.\n"
           "\n"
           "## Next steps\n"
           "What should happen next.\n"
           "\n"
           "## Files\n"
           "One bullet per file the agent read or edited, as `- <path>`.\n"
           "\n"
           "Rules:\n"
           "- Recall over precision: keep every fact the agent would otherwise "
           "have to rediscover; drop nothing load-bearing.\n"
           "- Be factual and terse; never invent details that are not in the "
           "transcript.\n"
           "- Discard raw tool output and verbatim file contents; summarize "
           "their substance instead.\n"
           "- Keep it under 800 words.";
}

/* ------------------------------------------------------------------ */
/* Config                                                              */
/* ------------------------------------------------------------------ */

void compact_config_from(unsigned long context_limit, double compact_threshold,
                         double clear_threshold, compact_config *out) {
    out->context_limit = context_limit;
    out->compact_threshold = compact_threshold;
    out->clear_threshold = clear_threshold;
}

bool compact_at(const compact_config *c, unsigned long *out) {
    if (c->compact_threshold <= 0.0)
        return false;
    *out = (unsigned long)(c->compact_threshold * (double)c->context_limit);
    return true;
}

bool clear_at(const compact_config *c, unsigned long *out) {
    if (c->clear_threshold <= 0.0)
        return false;
    *out = (unsigned long)(c->clear_threshold * (double)c->context_limit);
    return true;
}

/* ------------------------------------------------------------------ */
/* Token estimation                                                    */
/* ------------------------------------------------------------------ */

static size_t trc_bytes(const tool_result_content &c) {
    if (std::holds_alternative<trc_image>(c))
        return 8;
    if (const auto *t = std::get_if<trc_text>(&c))
        return t->text.size();
    if (const auto *s = std::get_if<trc_json>(&c))
        return s->value.size();
    return 0;
}

static size_t uc_bytes(const user_part &p) {
    if (const auto *t = std::get_if<user_text_part>(&p))
        return t->text.size();
    const auto &r = std::get<user_tool_result_part>(p);
    size_t total = r.name.size();
    for (const auto &c : r.content)
        total += trc_bytes(c);
    return total;
}

static size_t ac_bytes(const assistant_part &p) {
    if (std::holds_alternative<asst_image_part>(p))
        return 0;
    if (const auto *c = std::get_if<asst_tool_call_part>(&p))
        return c->name.size() + c->args.size();
    if (const auto *t = std::get_if<asst_text_part>(&p))
        return t->text.size();
    if (const auto *r = std::get_if<asst_reasoning_part>(&p))
        return r->text.size();
    return 0;
}

static size_t message_bytes(const message *m) {
    if (m->is_system())
        return m->system().text.size();
    size_t total = 0;
    if (m->is_user()) {
        for (const auto &p : m->user().parts)
            total += uc_bytes(p);
    } else if (m->is_assistant()) {
        for (const auto &p : m->assistant().parts)
            total += ac_bytes(p);
    }
    return total;
}

unsigned long compact_estimate_tokens(const message *msgs, size_t n) {
    unsigned long long bytes = 0;
    for (size_t i = 0; i < n; i++)
        bytes += message_bytes(&msgs[i]);
    return (unsigned long)((bytes + 3) / 4);
}

/* ------------------------------------------------------------------ */
/* Boundary + clearing                                                 */
/* ------------------------------------------------------------------ */

static bool compact_is_pure_text_user(const message *m) {
    return m->is_user() && m->user().parts.size() == 1 &&
           std::holds_alternative<user_text_part>(m->user().parts[0]);
}

size_t compact_keep_boundary(const message *msgs, size_t n) {
    for (size_t i = n; i-- > 0;) {
        if (compact_is_pure_text_user(&msgs[i]))
            return i;
    }
    return n;
}

size_t compact_clear_old_tool_results(message *msgs, size_t n, size_t boundary) {
    size_t cleared = 0;
    for (size_t i = 0; i < boundary && i < n; i++) {
        message *m = &msgs[i];
        if (!m->is_user())
            continue;
        for (auto &p : m->user().parts) {
            user_tool_result_part *r = std::get_if<user_tool_result_part>(&p);
            if (!r)
                continue;
            bool already = false;
            for (const auto &c : r->content) {
                const trc_text *t = std::get_if<trc_text>(&c);
                if (t && t->text == CLEARED_PLACEHOLDER)
                    already = true;
            }
            if (already)
                continue;
            r->content.clear();
            trc_text t;
            t.text = CLEARED_PLACEHOLDER;
            r->content.push_back(std::move(t));
            cleared++;
        }
    }
    return cleared;
}

/* ------------------------------------------------------------------ */
/* Transcripts                                                         */
/* ------------------------------------------------------------------ */

static void put_trc(xbuf *b, const tool_result_content &c) {
    if (std::holds_alternative<trc_image>(c)) {
        xbuf_put(b, "[image]");
        return;
    }
    if (const auto *t = std::get_if<trc_text>(&c)) {
        if (!t->text.empty())
            xbuf_put(b, t->text.c_str());
        return;
    }
    if (const auto *s = std::get_if<trc_json>(&c))
        if (!s->value.empty())
            xbuf_put(b, s->value.c_str());
}

char *compact_transcript(const message *msgs, size_t n) {
    xbuf b;
    xbuf_init(&b);
    for (size_t i = 0; i < n; i++) {
        const message *m = &msgs[i];
        if (m->is_system()) {
            xbuf_put(&b, "System: ");
            xbuf_put(&b, m->system().text.c_str());
            xbuf_push(&b, '\n');
        } else if (m->is_user()) {
            for (size_t j = 0; j < m->user().parts.size(); j++) {
                const auto &p = m->user().parts[j];
                if (const auto *t = std::get_if<user_text_part>(&p)) {
                    xbuf_put(&b, "User: ");
                    xbuf_put(&b, t->text.c_str());
                    xbuf_push(&b, '\n');
                } else {
                    const auto &r = std::get<user_tool_result_part>(p);
                    xbuf_fmt(&b, "ToolResult(%s): ", r.name.c_str());
                    for (const auto &c : r.content)
                        put_trc(&b, c);
                    xbuf_push(&b, '\n');
                }
            }
        } else if (m->is_assistant()) {
            for (size_t j = 0; j < m->assistant().parts.size(); j++) {
                const auto &p = m->assistant().parts[j];
                if (const auto *t = std::get_if<asst_text_part>(&p)) {
                    xbuf_put(&b, "Assistant: ");
                    xbuf_put(&b, t->text.c_str());
                    xbuf_push(&b, '\n');
                } else if (const auto *c = std::get_if<asst_tool_call_part>(&p)) {
                    xbuf_fmt(&b, "Assistant tool call %s: %s",
                             c->name.c_str(), c->args.c_str());
                    xbuf_push(&b, '\n');
                } else if (const auto *r = std::get_if<asst_reasoning_part>(&p)) {
                    xbuf_put(&b, "Assistant reasoning: ");
                    xbuf_put(&b, r->text.c_str());
                    xbuf_push(&b, '\n');
                }
            }
        }
    }
    return xbuf_steal(&b);
}

static void push_md_tool_result(xbuf *b, const tool_result_content &c) {
    if (const auto *t = std::get_if<trc_text>(&c)) {
        xbuf_put(b, t->text.c_str());
        xbuf_put(b, "\n\n");
    } else if (const auto *s = std::get_if<trc_json>(&c)) {
        xbuf_put(b, "```json\n");
        xbuf_put(b, s->value.c_str());
        xbuf_put(b, "\n```\n\n");
    } else {
        xbuf_put(b, "[image]\n\n");
    }
}

char *compact_transcript_markdown(const message *msgs, size_t n) {
    xbuf b;
    xbuf_init(&b);
    xbuf_put(&b, "# llsh transcript\n\n");
    if (n == 0) {
        xbuf_put(&b, "_No messages yet._\n");
        return xbuf_steal(&b);
    }
    for (size_t i = 0; i < n; i++) {
        const message *m = &msgs[i];
        if (m->is_system()) {
            xbuf_put(&b, "## System\n\n");
            xbuf_put(&b, m->system().text.c_str());
            xbuf_put(&b, "\n\n");
        } else if (m->is_user()) {
            for (size_t j = 0; j < m->user().parts.size(); j++) {
                const auto &p = m->user().parts[j];
                if (const auto *t = std::get_if<user_text_part>(&p)) {
                    xbuf_put(&b, "## User\n\n");
                    xbuf_put(&b, t->text.c_str());
                    xbuf_put(&b, "\n\n");
                } else {
                    const auto &r = std::get<user_tool_result_part>(p);
                    xbuf_fmt(&b, "### Tool result: `%s`\n\n",
                             r.name.c_str());
                    for (const auto &c : r.content)
                        push_md_tool_result(&b, c);
                }
            }
        } else if (m->is_assistant()) {
            for (size_t j = 0; j < m->assistant().parts.size(); j++) {
                const auto &p = m->assistant().parts[j];
                if (const auto *t = std::get_if<asst_text_part>(&p)) {
                    xbuf_put(&b, "## Assistant\n\n");
                    xbuf_put(&b, t->text.c_str());
                    xbuf_put(&b, "\n\n");
                } else if (const auto *c = std::get_if<asst_tool_call_part>(&p)) {
                    xbuf_fmt(&b, "### Tool call: `%s`\n\n",
                             c->name.c_str());
                    if (c->args.empty() || c->args == "null") {
                        xbuf_put(&b, "_no arguments_\n\n");
                    } else {
                        xbuf_put(&b, "```json\n");
                        xbuf_put(&b, c->args.c_str());
                        xbuf_put(&b, "\n```\n\n");
                    }
                } else if (const auto *r = std::get_if<asst_reasoning_part>(&p)) {
                    if (!r->text.empty()) {
                        xbuf_put(&b, "> **reasoning** ");
                        xbuf_put(&b, r->text.c_str());
                        xbuf_put(&b, "\n\n");
                    }
                }
            }
        }
    }
    return xbuf_steal(&b);
}

/* ------------------------------------------------------------------ */
/* File-path extraction                                                */
/* ------------------------------------------------------------------ */

static bool starts_with_i(const char *s, size_t n, const char *prefix) {
    size_t pn = strlen(prefix);
    if (n < pn)
        return false;
    for (size_t i = 0; i < pn; i++) {
        char a = s[i];
        if (a >= 'A' && a <= 'Z')
            a = (char)(a - 'A' + 'a');
        if (a != prefix[i])
            return false;
    }
    return true;
}

void compact_recent_file_paths(const char *summary, strlist *out) {
    sl_init(out);
    if (!summary)
        return;
    bool in_files = false;
    const char *p = summary;
    while (*p) {
        const char *nl = strchr(p, '\n');
        size_t linelen = nl ? (size_t)(nl - p) : strlen(p);
        const char *line_end = p + linelen;
        size_t tlen;
        const char *line = xtrim_span(p, linelen, &tlen);
        if (tlen > 0 && line[0] == '#') {
            const char *h = line;
            while (*h == '#')
                h++;
            size_t hl;
            h = xtrim_span(h, (size_t)(line_end - h), &hl);
            in_files = starts_with_i(h, hl, "files");
        } else if (in_files && tlen > 0 && (line[0] == '-' || line[0] == '*')) {
            const char *rest = line + 1;
            size_t rl;
            rest = xtrim_span(rest, (size_t)(line_end - rest), &rl);
            if (rl > 0)
                sl_push(out, xstrndup(rest, rl));
        }
        p = nl ? nl + 1 : p + linelen;
    }
}

/* ------------------------------------------------------------------ */
/* Compose history                                                     */
/* ------------------------------------------------------------------ */

history *compact_compose_history(const char *summary,
                                  const char **file_paths,
                                  const char **file_contents, size_t nfiles,
                                   const message *tail, size_t ntail) {
    xbuf body;
    xbuf_init(&body);
    xbuf_put(&body,
             "The earlier conversation was summarized to fit the context "
             "window. Continue from this summary.\n\n");
    if (summary) {
        size_t tlen;
        const char *s = xtrim_span(summary, strlen(summary), &tlen);
        xbuf_putn(&body, s, tlen);
        xbuf_push(&body, '\n');
    }

    unsigned long total = 0;
    size_t taken = 0;
    for (size_t i = 0; i < nfiles && i < MAX_FILES; i++) {
        if (total >= FILES_CAP)
            break;
        const char *content = file_contents[i] ? file_contents[i] : "";
        size_t cpoints = utf8_len(content);
        unsigned long take = PER_FILE_CAP;
        if (take > FILES_CAP - total)
            take = FILES_CAP - total;
        if (take > cpoints)
            take = cpoints;
        if (take == 0)
            break;
        size_t take_bytes = utf8_byte_len(content, take);
        total += take;
        taken++;
        xbuf_fmt(&body, "\n### %s\n```\n", file_paths[i] ? file_paths[i] : "");
        xbuf_putn(&body, content, take_bytes);
        xbuf_put(&body, "\n```\n");
    }

    history *h = (history *) xmalloc(sizeof(*h));
    new (h) history{};
    char *body_text = xbuf_steal(&body);
    h->msgs.push_back(msg_user_text(body_text));
    free(body_text);
    for (size_t i = 0; i < ntail; i++)
        h->msgs.push_back(tail[i]);
    (void)taken;
    return h;
}

/* ------------------------------------------------------------------ */
/* SelfCompact-style rubric probe (arXiv:2606.23525)                   */
/* ------------------------------------------------------------------ */

const char *compact_rubric_preamble(void) {
    return "You are a careful judge for a terminal coding agent. You are "
           "given the agent's trajectory and a search rubric of four yes/no "
           "questions (C1, C2, C3, N1). Answer each question exactly in the "
           "requested line format, with the evidence it asks for. Do not "
           "summarize the trajectory, do not restate the questions, and add "
           "no text beyond the four verdict lines.";
}

const char *compact_rubric_prompt(void) {
    return "You are about to decide whether to compress your conversation "
           "history into a summary that REPLACES the full history above. "
           "After compression, you continue from only [system, original "
           "question, assistant summary, user continue]. Compression is "
           "irreversible: anything not preserved in the summary is gone.\n"
           "\n"
           "Compression is safe ONLY when ALL FOUR of the following hold: "
           "(C1) the trajectory has reached a closed unit (not mid-thought), "
           "(C2) the essential information is reducible to 3-5 citable "
           "facts without loss, (C3) something has progressed since the last "
           "compression, (N1) you are NOT currently stuck in a way "
           "summarization would mask. Answer C1, C2, C3, N1 honestly. Each "
           "Y answer requires verbatim evidence quoted from the trajectory "
           "above; answers without evidence default to N.\n"
           "\n"
           "C1 CLOSED-UNIT: The most recent assistant message is a closed "
           "unit - a completed tool call whose result is now visible, or a "
           "completed sub-analysis with a clear stopping point. It is NOT "
           "mid-sentence reasoning (\"Let me now check...\", \"I should "
           "next look at...\"), and not a half-formulated query. *If Y*, "
           "quote the closing fragment of the last assistant message. *If "
           "N*, quote the open fragment that shows the trajectory is "
           "mid-thought.\n"
           "C2 SUMMARIZABLE: You can write 3-5 essential facts (with "
           "verbatim citations from the trajectory) that future-you needs "
           "to continue after compression. Each fact must be a single "
           "concrete statement: a name, date, URL, quoted claim, or "
           "resolved sub-question. Answer N if the trajectory's value is "
           "dispersed across many small inferences (e.g., a list of dead-"
           "end queries needed to avoid retries, negative results that "
           "constrain the hypothesis space) that would be lost without the "
           "dispersal. *If Y*, list the 3-5 facts numbered, each with a "
           "verbatim citation in quotes. *If N*, name in one sentence the "
           "class of information that would be lost.\n"
           "C3 PROGRESS: Since the most recent compression (or since the "
           "start of the conversation if none), you have either obtained a "
           "new concrete fact (name, date, URL, or quoted claim) OR "
           "refined the sub-question being pursued. *If Y*, name the new "
           "fact or refined sub-question. *If N*, state that you are "
           "returning the same state you compressed from.\n"
           "N1 STUCK: At least 3 of your last 4 tool calls returned no new "
           "information (i.e., were duplicates or returned already-known "
           "content). If you have made fewer than 4 tool calls total, "
           "answer N. *If Y*, name 1 distinct strategy you have NOT yet "
           "tried (different tool, different query type, different angle "
           "on the question). *If N*, name one new fact obtained "
           "recently.\n"
           "\n"
           "Output: exactly 4 lines, no preamble or trailing text:\n"
           "C1: Y/N -- <evidence>\n"
           "C2: Y/N -- <evidence>\n"
           "C3: Y/N -- <evidence>\n"
           "N1: Y/N -- <evidence>";
}

/* First line whose start matches `label` (case-insensitive, followed by
 * whitespace or ':'), or NULL. */
static const char *rubric_line(const char *answer, const char *label) {
    size_t ll = strlen(label);
    const char *p = answer;
    while (*p) {
        const char *nl = strchr(p, '\n');
        size_t len = nl ? (size_t)(nl - p) : strlen(p);
        if (len >= ll && strncasecmp(p, label, ll) == 0 &&
            (len == ll || p[ll] == ' ' || p[ll] == '\t' || p[ll] == ':'))
            return p;
        if (!nl)
            break;
        p = nl + 1;
    }
    return NULL;
}

/* Verdict at `line` (just past the label). Sets *evidence to the
 * remainder of the line after an optional `--` separator. Returns 1 = Y,
 * 0 = N, -1 = no verdict token. A Y with empty evidence is returned as N
 * (the paper's "answers without evidence default to N"). */
static int rubric_verdict(const char *line, const char **evidence) {
    while (*line == ' ' || *line == '\t')
        line++;
    if (*line == ':') {
        line++;
        while (*line == ' ' || *line == '\t')
            line++;
    }
    char v = (char)toupper((unsigned char)*line);
    if (v != 'Y' && v != 'N')
        return -1;
    int verdict = (v == 'Y');
    line++;
    while (*line == ' ' || *line == '\t')
        line++;
    if (line[0] == '-' && line[1] == '-')
        line += 2;
    while (*line == ' ' || *line == '\t')
        line++;
    *evidence = line;
    if (verdict) {
        /* A Y needs real evidence on the same line (the paper's rule). */
        const char *p = line;
        while (*p && *p != '\n')
            p++;
        if (p == line)
            verdict = 0;
    }
    return verdict;
}

int compact_rubric_parse(const char *answer, compact_rubric *out,
                         char **err) {
    if (err)
        *err = NULL;
    if (!out)
        return -1;
    out->c1 = out->c2 = out->c3 = out->n1 = 0;
    if (!answer || !*answer) {
        if (err)
            *err = xstrdup("rubric: empty verdict");
        return -1;
    }
    struct {
        const char *label;
        int *slot;
    } items[4] = {
        {"C1", &out->c1},
        {"C2", &out->c2},
        {"C3", &out->c3},
        {"N1", &out->n1},
    };
    for (int i = 0; i < 4; i++) {
        const char *line = rubric_line(answer, items[i].label);
        if (!line) {
            if (err)
                *err = xstrfmt("rubric: missing `%s` verdict line",
                               items[i].label);
            return -1;
        }
        const char *ev = NULL;
        int verdict = rubric_verdict(line + strlen(items[i].label), &ev);
        if (verdict < 0) {
            if (err)
                *err = xstrfmt("rubric: malformed `%s` verdict line",
                               items[i].label);
            return -1;
        }
        *items[i].slot = verdict;
    }
    return 0;
}

bool compact_rubric_fire(const compact_rubric *r) {
    return r && r->c1 && r->c2 && r->c3 && !r->n1;
}

bool compact_rubric_due(unsigned long iteration, unsigned long prompt_tokens,
                        unsigned long summaries_so_far,
                        unsigned long max_summaries,
                        unsigned long rounds_since_probe) {
    if (iteration < RUBRIC_MIN_ITERATION)
        return false;
    if (prompt_tokens < RUBRIC_MIN_TOKENS)
        return false;
    if (max_summaries > 0 && summaries_so_far >= max_summaries)
        return false;
    if (rounds_since_probe < RUBRIC_MIN_ROUNDS_SINCE_PROBE)
        return false;
    return true;
}