/* context.c — Hot-path context assembly for small-model agent framework.
 *
 *   The small model has a fixed context budget (default 3072 tokens). This
 *   module assembles a prompt from tiered priority sources with strict budgets.
 *   Everything is token-counted; nothing overflows. Built as a shared lib
 *   callable from Python via ctypes.
 *
 *   Pri0:  System + tool specs  (512 tok, immutable, cron-updated)
 *   Pri1:  Active task          (768 tok, per-task reset)
 *   Pri2:  Persistent memory    (1024 tok, scored + deduplicated)
 *   Pri3:  Recent conversation  (768 tok, FIFO sliding window)
 *
 *   Compile:
 *     cc -O3 -std=c11 -fPIC -shared context.c -o libcontext.so
 *
 *   Token estimation: ~4 chars/token for English text (conservative).
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#define MAX_TOKENS   3072
#define PRI0_TOKENS  512
#define PRI1_TOKENS  768
#define PRI2_TOKENS  1024
#define PRI3_TOKENS  768
#define CHARS_PER_TOKEN 4

#define MAX_BUF (MAX_TOKENS * CHARS_PER_TOKEN * 2)  /* generous */

static char g_buf[MAX_BUF];

/* ---- Token estimation ---- */
static int est_tokens(const char *s) {
    int chars = 0;
    while (*s) { chars++; s++; }
    return (chars + CHARS_PER_TOKEN - 1) / CHARS_PER_TOKEN;
}

/* Truncate string to fit max_tokens, adding "..." if truncated.
 * Returns bytes written (not counting null). */
static int truncate(char *dst, const char *src, int max_tokens) {
    int max_chars = max_tokens * CHARS_PER_TOKEN;
    int n = 0;
    while (src[n] && n < max_chars - 3) n++;
    memcpy(dst, src, n);
    if (src[n]) {
        dst[n++] = '.'; dst[n++] = '.'; dst[n++] = '.';
    }
    dst[n] = '\0';
    return n;
}

/* ---- Public API ---- */

/* Score a memory entry by novelty. Returns 0.0-1.0.
 * Simple: count word overlap with existing entries. Lower overlap = higher novelty. */
double context_memory_score(const char *entry, const char **existing, int n_existing) {
    if (n_existing == 0) return 1.0;
    int total_words = 0, overlap = 0;
    const char *p = entry;
    char word[64];
    while (*p) {
        while (*p == ' ' || *p == '\n') p++;
        if (!*p) break;
        int len = 0;
        while (p[len] && p[len] != ' ' && p[len] != '\n' && len < 63) len++;
        memcpy(word, p, len); word[len] = '\0';
        p += len;
        total_words++;
        for (int i = 0; i < n_existing; i++) {
            const char *found = strstr(existing[i], word);
            if (found) {
                int at_start = (found == existing[i]);
                int at_boundary = at_start || (found[-1] == ' ' || found[-1] == '\n');
                if (at_boundary) { overlap++; break; }
            }
        }
    }
    if (total_words == 0) return 0.0;
    return 1.0 - (double)overlap / (double)total_words;
}

/* Build the full context buffer from tiered sources.
 *
 *   pri0:  system prompt
 *   pri1:  current task description
 *   pri2:  array of memory entries (pre-sorted by score)
 *   npri2: number of memory entries
 *   pri3:  array of recent messages (newest last)
 *   npri3: number of recent messages
 *
 * Returns the assembled prompt string. Static buffer — copy if needed. */
const char *context_assemble(
    const char *pri0,
    const char *pri1,
    const char **pri2, int npri2,
    const char **pri3, int npri3)
{
    char *p = g_buf;
    int remaining = MAX_TOKENS;

    /* Pri0: system prompt — always included, truncated if needed */
    int tok = est_tokens(pri0);
    if (tok > PRI0_TOKENS) tok = PRI0_TOKENS;
    truncate(p, pri0, tok);
    p += strlen(p);
    *p++ = '\n'; *p = '\0';
    remaining -= tok;

    /* Pri1: task — always included */
    tok = est_tokens(pri1);
    if (tok > PRI1_TOKENS) tok = PRI1_TOKENS;
    if (tok <= remaining) {
        truncate(p, pri1, tok);
        p += strlen(p);
        *p++ = '\n'; *p = '\0';
        remaining -= tok;
    }

    /* Pri2: scored memories — fill remaining up to PRI2_TOKENS */
    int pri2_budget = PRI2_TOKENS;
    if (pri2_budget > remaining) pri2_budget = remaining;
    if (pri2_budget > 0 && npri2 > 0) {
        p += sprintf(p, "--- RELEVANT MEMORIES ---\n");
        for (int i = 0; i < npri2 && pri2_budget > 0; i++) {
            int len = strlen(pri2[i]);
            int et = est_tokens(pri2[i]);
            if (et > pri2_budget) et = pri2_budget;
            if (et <= 0) continue;
            truncate(p, pri2[i], et);
            int wrote = strlen(p);
            p += wrote;
            *p++ = '\n'; *p = '\0';
            pri2_budget -= et;
        }
        *p++ = '\n'; *p = '\0';
        remaining -= (PRI2_TOKENS - pri2_budget);
    }

    /* Pri3: recent messages — fill remaining */
    if (remaining > 0 && npri3 > 0) {
        /* Include from newest backward */
        int start = npri3 - 1;
        int budget = PRI3_TOKENS;
        if (budget > remaining) budget = remaining;
        /* Find how many messages fit */
        int msgs_to_include = 0;
        int used = 0;
        for (int i = start; i >= 0; i--) {
            int et = est_tokens(pri3[i]);
            if (used + et > budget) break;
            msgs_to_include++;
            used += et;
        }
        /* Write them oldest-first */
        for (int i = start - msgs_to_include + 1; i <= start; i++) {
            if (i >= 0) {
                int len = strlen(pri3[i]);
                memcpy(p, pri3[i], len);
                p += len;
                *p++ = '\n'; *p = '\0';
            }
        }
        remaining -= used;
    }

    return g_buf;
}

/* Compress tool output to gist. Returns static buffer.
 * Simple heuristic version (full compression uses LLM pass). */
const char *context_compress_output(const char *raw, int max_tokens) {
    char *p = g_buf;
    /* Count lines and total chars */
    int lines = 0, chars = 0;
    const char *s;
    for (s = raw; *s; s++) {
        chars++;
        if (*s == '\n') lines++;
    }
    /* Extract first and last few lines + stats */
    int budget = max_tokens * CHARS_PER_TOKEN;
    p += snprintf(p, budget, "[%d lines, %d chars]\n", lines, chars);

    /* First 2 non-empty lines */
    const char *line_start = raw;
    int lc = 0;
    for (s = raw; *s && lc < 4; s++) {
        if (*s == '\n' || *s == '\0') {
            if (s > line_start && p - g_buf < budget - 20) {
                int n = s - line_start;
                if (n > 80) n = 80;  /* max 80 chars per line */
                memcpy(p, line_start, n);
                p += n;
                if (s - line_start > 80) { *p++ = '.'; *p++ = '.'; *p++ = '.'; }
                *p++ = '\n';
                lc++;
            }
            line_start = s + 1;
            if (*s == '\0') break;
        }
    }
    /* Last line if there's space */
    if (lines > 4 && p - g_buf < budget - 40) {
        p += snprintf(p, 40, "...\n");
        /* Find last non-empty line */
        const char *last = raw + strlen(raw) - 1;
        while (last > raw && *last == '\n') last--;
        const char *ll = last;
        while (ll > raw && ll[-1] != '\n') ll--;
        int n = last - ll + 1;
        if (n > 80) { ll = last - 79; n = 80; }
        memcpy(p, ll, n);
        p += n;
    }
    *p = '\0';
    return g_buf;
}

/* Quick response format check — does it match the THOUGHT/ACTION template? */
int context_validate_response(const char *response) {
    return (strstr(response, "THOUGHT:") != NULL ||
            strstr(response, "ACTION:")  != NULL);
}