/*
 * kcpp / ICPP - the Internet C++ compiler and interpreter
 *
 * Copyright (c) 2026 Muhammad Anisur Rahman
 * Free to use and distribute under the MIT License.
 *
 * SPDX-License-Identifier: MIT
 * See the LICENSE file at the repository root for full terms.
 */

/*---------------------------------------------------------------------------
  icpp_llm.c - the LLM lifecycle API for ICPP scripts.

  Pull in with:   #include "icpp_llm.c"

  This is the layer that was missing. The tensor engine, the autodiff rules
  and the chatbot retrieval were all working, but every script had to
  hand-build its own model, wire its own backward chain, and invent its own
  file names. This library turns all of that into:

        create_llm("models", "support");     <- once, ever
        llm_open("models", "support");
        llm_train_file("faq.txt", 200, 0.05);
        llm_save();
        iprint(llm_answer("how do I reset my password", 80));

  MODEL ON DISK. create_llm() lays down one directory per model, so a model
  is a thing you can copy, back up, or ship:

        <dirname>/<llmname>_emb.t     embedding table   [256, dim]
        <dirname>/<llmname>_w.t       LM head weights   [dim, 256]
        <dirname>/<llmname>_b.t       LM head bias      [1, 256]
        <dirname>/<llmname>_kb.txt    knowledge base, one fact per line
        <dirname>/<llmname>.meta      dim, vocab, version

  ARCHITECTURE. Byte-level next-token prediction: a 256-entry embedding
  table feeding a linear head back to 256 logits. Small on purpose - it
  trains in seconds inside the interpreter, and for a support bot the
  retrieval leg carries the facts while the model supplies phrasing. The
  same file layout extends to a deeper stack later without changing this
  API, because callers never touch the tensors directly.

  THE TWO LEGS OF AN ANSWER. llm_answer() retrieves before it generates:
  the knowledge base is searched by cosine similarity for the fact that
  best matches the question, and that fact is what gets returned as the
  substance. Generation supplies continuation text. Retrieval is what makes
  the answer TRUE; generation only makes it read like prose. A bot that
  generates without retrieving is a bot that invents policy.

  USING IT FROM A WEB PAGE. iwebserver serves <?icpp ?> pages, so a support
  page is just:

        <?icpp
          #include "icpp_llm.c"
          $q = icpp_getparam("q");
          llm_open("models", "support");
          iprint("<p>", llm_answer($q, 120), "</p>");
          llm_close();
        ?>

  Open per request and close after: models are files, and reopening one is
  three tensor_load calls.
---------------------------------------------------------------------------*/

/* The backward rules this library trains with - tensor_linear_backward_w,
   _x and _b - are ICPP functions, not builtins. Without this include the
   calls resolved to nothing and training hung silently rather than
   reporting an undefined function. */
#include "icpp_autodiff.c"

/* ---- loaded-model state -------------------------------------------- */

int g_llm_emb;          /* [256, dim] embedding table, 0 when closed  */
int g_llm_w;            /* [dim, 256] head weights                    */
int g_llm_b;            /* [1, 256]   head bias                       */
int g_llm_dim;          /* embedding width of the open model          */
int g_llm_open;         /* 1 when a model is loaded                   */
/* char*, NOT int. These hold STRINGS, and a string variable declared int
   never carries token==STRING - so every builtin that consumes one sees a
   number instead. Declared int, llm_file() built paths out of raw pointer
   values ("561987411/561987418_kb.txt"), the knowledge base was written
   somewhere unreachable, and llm_info() printed the model's name as a
   number. Second time this exact trap has bitten in this file. */
char *g_llm_dir;        /* directory of the open model                */
char *g_llm_name;       /* name of the open model                     */

#define LLM_VOCAB   256
#define LLM_DEFDIM   32

/* ---- paths ---------------------------------------------------------- */

function llm_file(dirname, llmname, suffix)
{
    return isprint(dirname, "/", llmname, suffix);
}

/* ---- deterministic small init --------------------------------------
   Weights must not start identical or every unit learns the same thing.
   A fixed arithmetic spread is used rather than a random source so that
   creating the same model twice gives the same starting point - a model
   you can reproduce is a model you can debug. */

function llm_init_spread(rows, cols, seed)
{
    int t;
    t = tensor_new(rows, cols);
    /* One builtin call rather than rows*cols interpreted iterations.
       Doing this in a script loop made create_llm() take minutes: a
       256x32 table plus a 32x256 head is 16,384 iterations, and the
       interpreted loop, not the arithmetic, was the whole cost. */
    tensor_init_spread(t, seed);
    return t;
}

/* ---- creation -------------------------------------------------------- */

/* Create a new model with an explicit embedding width. Writes every file
   the model needs and leaves nothing loaded - call llm_open() next.
   Returns 1 on success, 0 if the model already exists. */
function create_llm_dim(dirname, llmname, dim)
{
    int emb;
    int w;
    int b;
    char *meta;         /* a PATH - must be char*, not int: every
                           string-consuming builtin dispatches on
                           token==STRING, and an int-declared variable
                           never carries that tag no matter what pointer
                           is stored in it. Declaring these int meant
                           ifprint() silently wrote nothing. */

    icpp_mkdir(dirname);

    meta = llm_file(dirname, llmname, ".meta");
    if (icpp_file_exists(meta) == 1)
    {
        iprint("create_llm: '", llmname, "' already exists in ", dirname);
        return 0;
    }

    emb = llm_init_spread(LLM_VOCAB, dim, 9);
    w   = llm_init_spread(dim, LLM_VOCAB, 10);
    b   = tensor_new(1, LLM_VOCAB);

    tensor_save(emb, llm_file(dirname, llmname, "_emb.t"));
    tensor_save(w,   llm_file(dirname, llmname, "_w.t"));
    tensor_save(b,   llm_file(dirname, llmname, "_b.t"));

    tensor_free(emb);
    tensor_free(w);
    tensor_free(b);

    icpp_remove(meta);
    ifprint(meta, "llm_version 1");
    ifprint(meta, "name ", llmname);
    ifprint(meta, "dim ", dim);
    ifprint(meta, "vocab ", LLM_VOCAB);

    /* An empty knowledge base, so llm_kb_* works before anything is added. */
    icpp_remove(llm_file(dirname, llmname, "_kb.txt"));

    iprint("created llm '", llmname, "' in ", dirname,
           " (dim ", dim, ", vocab ", LLM_VOCAB, ")");
    return 1;
}

/* Create a model with the default width. This is the call to reach for. */
function create_llm(dirname, llmname)
{
    return create_llm_dim(dirname, llmname, LLM_DEFDIM);
}

/* ---- open / save / close --------------------------------------------- */

function llm_open(dirname, llmname)
{
    char *meta;
    int fh;
    int parts;
    char *line;

    if (g_llm_open == 1)
        llm_close();

    meta = llm_file(dirname, llmname, ".meta");
    if (icpp_file_exists(meta) == 0)
    {
        iprint("llm_open: no model '", llmname, "' in ", dirname,
               " - call create_llm first");
        return 0;
    }

    /* Recover dim from the meta file rather than assuming the default,
       so models of different widths can coexist. */
    g_llm_dim = LLM_DEFDIM;
    fh = icpp_fopen(meta);
    if (fh != 0)
    {
        while (icpp_feof(fh) == 0)
        {
            line = icpp_fgets(fh);
            parts = icpp_split(icpp_trim(line), " ");
            if (icpp_list_count(parts) > 1)
            {
                if (icpp_strcmp(icpp_list_get(parts, 0), "dim") == 0)
                    g_llm_dim = icpp_atoi(icpp_list_get(parts, 1));
            }
            icpp_list_free(parts);
        }
        icpp_fclose(fh);
    }

    g_llm_emb = tensor_load(llm_file(dirname, llmname, "_emb.t"));
    g_llm_w   = tensor_load(llm_file(dirname, llmname, "_w.t"));
    g_llm_b   = tensor_load(llm_file(dirname, llmname, "_b.t"));

    g_llm_dir = dirname;
    g_llm_name = llmname;
    g_llm_open = 1;
    return 1;
}

function llm_save()
{
    if (g_llm_open == 0)
    {
        iprint("llm_save: no model open");
        return 0;
    }
    tensor_save(g_llm_emb, llm_file(g_llm_dir, g_llm_name, "_emb.t"));
    tensor_save(g_llm_w,   llm_file(g_llm_dir, g_llm_name, "_w.t"));
    tensor_save(g_llm_b,   llm_file(g_llm_dir, g_llm_name, "_b.t"));
    return 1;
}

function llm_close()
{
    if (g_llm_open == 0)
        return 0;
    tensor_free(g_llm_emb);
    tensor_free(g_llm_w);
    tensor_free(g_llm_b);
    g_llm_emb = 0;
    g_llm_w = 0;
    g_llm_b = 0;
    g_llm_open = 0;
    return 1;
}

function llm_info()
{
    if (g_llm_open == 0)
    {
        iprint("llm: no model open");
        return 0;
    }
    iprint("llm '", g_llm_name, "' in ", g_llm_dir,
           ": dim ", g_llm_dim, ", vocab ", LLM_VOCAB,
           ", kb facts ", llm_kb_count());
    return 1;
}

/* ---- training --------------------------------------------------------
   One gradient step over a whole string, repeated 'steps' times. The
   target of position i is the byte at position i+1, which is what makes
   this next-token prediction. Returns the final loss so a caller can see
   whether it actually learned - a training call that cannot report loss
   is a training call you cannot trust. */

function llm_train_text(text, steps, lr)
{
    int ids;
    int targets;
    int n;
    int i;
    int step;
    int x;
    int logits;
    int dlogits;
    int dw;
    int db;
    int dx;
    int demb;
    double loss;

    if (g_llm_open == 0)
    {
        iprint("llm_train_text: no model open");
        return 0.0;
    }

    ids = tensor_tokenize(text);
    n = tensor_cols(ids);
    if (n < 2)
    {
        tensor_free(ids);
        iprint("llm_train_text: need at least 2 bytes of text");
        return 0.0;
    }

    /* Shift by one: predict the next byte. The final byte has no target,
       so the usable length is n-1. */
    targets = tensor_new(1, n - 1);
    for (i = 0; i < n - 1; i = i + 1)
        tensor_set(targets, 0, i, tensor_get(ids, 0, i + 1));

    /* Inputs are the first n-1 bytes. */
    x = tensor_new(1, n - 1);
    for (i = 0; i < n - 1; i = i + 1)
        tensor_set(x, 0, i, tensor_get(ids, 0, i));
    tensor_free(ids);
    ids = x;

    loss = 0.0;
    for (step = 0; step < steps; step = step + 1)
    {
        x = tensor_embedding_lookup(g_llm_emb, ids);
        logits = tensor_linear(x, g_llm_w, g_llm_b);
        loss = tensor_cross_entropy_loss(logits, targets);

        dlogits = tensor_cross_entropy_loss_backward(logits, targets);
        dw = tensor_linear_backward_w(x, g_llm_w, dlogits);
        db = tensor_linear_backward_b(dlogits);
        dx = tensor_linear_backward_x(x, g_llm_w, dlogits);
        demb = tensor_embedding_lookup_backward(g_llm_emb, ids, dx);

        tensor_sgd_step(g_llm_w, dw, lr);
        tensor_sgd_step(g_llm_b, db, lr);
        tensor_sgd_step(g_llm_emb, demb, lr);

        tensor_free(x);
        tensor_free(logits);
        tensor_free(dlogits);
        tensor_free(dw);
        tensor_free(db);
        tensor_free(dx);
        tensor_free(demb);
    }

    tensor_free(ids);
    tensor_free(targets);
    return loss;
}

/* Train on every line of a file - the usual way to feed an FAQ. */
function llm_train_file(path, steps, lr)
{
    int n;
    int i;
    double loss;

    n = icpp_file_lines(path);
    if (n == 0)
    {
        iprint("llm_train_file: ", path, " is empty or missing");
        return 0.0;
    }
    loss = 0.0;
    for (i = 1; i <= n; i = i + 1)
        loss = llm_train_text(icpp_file_line(path, i), steps, lr);
    return loss;
}

/* ---- generation ------------------------------------------------------ */

/* Continue from 'seedtext' for 'count' bytes. Returns the generated text
   only, not the seed. */
function llm_generate(seedtext, count)
{
    int ids;
    int out;
    int cur;
    int i;
    int one;
    int x;
    int logits;

    if (g_llm_open == 0)
    {
        iprint("llm_generate: no model open");
        return "";
    }
    if (count < 1)
        return "";

    ids = tensor_tokenize(seedtext);
    if (tensor_cols(ids) < 1)
    {
        tensor_free(ids);
        return "";
    }
    cur = tensor_get(ids, 0, tensor_cols(ids) - 1);
    tensor_free(ids);

    out = tensor_new(1, count);
    one = tensor_new(1, 1);
    for (i = 0; i < count; i = i + 1)
    {
        tensor_set(one, 0, 0, cur);
        x = tensor_embedding_lookup(g_llm_emb, one);
        logits = tensor_linear(x, g_llm_w, g_llm_b);
        cur = tensor_argmax(logits, 0);
        tensor_set(out, 0, i, cur);
        tensor_free(x);
        tensor_free(logits);
    }
    tensor_free(one);
    return tensor_detokenize(out);
}

/* ---- knowledge base (retrieval) -------------------------------------- */

function llm_kb_path()
{
    return llm_file(g_llm_dir, g_llm_name, "_kb.txt");
}

function llm_kb_add(fact)
{
    if (g_llm_open == 0)
    {
        iprint("llm_kb_add: no model open");
        return 0;
    }
    ifprint(llm_kb_path(), fact);
    return 1;
}

/* Add a fact together with the words a user is likely to type.
   NOT a shortcut - it is the fix for the problem that sinks pure
   text matching. "when are you open" shares no meaningful word with
   "Our support hours are 9am to 6pm", and "money back" shares none with
   "Refunds are issued within 14 days". No scoring tweak recovers a word
   that is not there. Real support systems solve this the same way: the
   answer carries the vocabulary its askers use.
   Stored as  <fact> |kw| <keywords>  on one line; only the fact is ever
   shown to the user. */
function llm_kb_add_kw(fact, keywords)
{
    if (g_llm_open == 0)
    {
        iprint("llm_kb_add_kw: no model open");
        return 0;
    }
    ifprint(llm_kb_path(), isprint(fact, " |kw| ", keywords));
    return 1;
}

/* The visible part of a stored line - everything before the keywords. */
function llm_kb_fact(line)
{
    int at;
    at = icpp_strstr(line, " |kw| ");
    if (at < 0)
        return line;
    return icpp_substr(line, 0, at);
}

function llm_kb_count()
{
    if (g_llm_open == 0)
        return 0;
    return icpp_file_lines(llm_kb_path());
}

/* The 1-based line of the fact closest to 'question' by cosine similarity
   of bigram embeddings. 0 when the knowledge base is empty.
   Embedding width is 256 by convention - 64 collides too often on real
   sentences, which showed up as a mis-ranked retrieval. */
/* The 1-based line of the fact that best matches 'question'.
   0 when the knowledge base is empty.

   SCORED TWO WAYS, because neither alone is good enough:

     word overlap  - how many of the question's meaningful words appear
                     in the fact, as a fraction of the question. This is
                     what actually distinguishes "when are you open" from
                     a fact about firmware updates.
     bigram cosine - icpp_text_embed similarity, which catches wording
                     that overlaps without matching exactly.

   Cosine alone was tried first and got 2 of 4 test questions wrong: it
   matched "can I get my money back" to a fact about a blinking LED,
   because character bigrams say little about meaning. Overlap dominates
   the score and cosine breaks ties.

   Words of 3 characters or fewer are ignored - "how", "do", "I", "my",
   "the" appear in everything and only add noise. */
function llm_kb_best(question)
{
    int n;
    int i;
    int j;
    int k;
    int best;
    int qwords;
    int lwords;
    int qn;
    int ln;
    int hits;
    int meaningful;
    int qvec;
    int lv;
    double score;
    double bestscore;
    char *qw;
    char *line;

    n = llm_kb_count();
    if (n == 0)
        return 0;

    qwords = icpp_split(icpp_lower(question), " ");
    qn = icpp_list_count(qwords);

    qvec = icpp_text_embed(question, 256);
    best = 0;
    bestscore = 0.0 - 2.0;

    for (i = 1; i <= n; i = i + 1)
    {
        /* Match against the whole stored line: the fact AND its
           keywords. The keywords exist precisely to be matched. */
        line = icpp_file_line(llm_kb_path(), i);
        lwords = icpp_split(icpp_lower(line), " ");
        ln = icpp_list_count(lwords);

        hits = 0;
        meaningful = 0;
        for (j = 0; j < qn; j = j + 1)
        {
            qw = icpp_list_get(qwords, j);
            if (icpp_strlen(qw) > 3)
            {
                meaningful = meaningful + 1;
                for (k = 0; k < ln; k = k + 1)
                {
                    /* strstr rather than equality so "opens"/"open" and
                       "refunds"/"refund" still count. */
                    if (icpp_strstr(icpp_list_get(lwords, k), qw) >= 0)
                    {
                        hits = hits + 1;
                        k = ln;
                    }
                }
            }
        }
        icpp_list_free(lwords);

        lv = icpp_text_embed(line, 256);
        score = tensor_cosine(qvec, lv) * 0.25;
        tensor_free(lv);
        if (meaningful > 0)
            score = score + ((double) hits / (double) meaningful);

        if (score > bestscore)
        {
            bestscore = score;
            best = i;
        }
    }

    icpp_list_free(qwords);
    tensor_free(qvec);
    return best;
}

/* ---- the answer ------------------------------------------------------
   Retrieval first, generation second. The retrieved fact is the answer's
   substance; the generated tail is only phrasing. When the knowledge base
   is empty this says so rather than inventing an answer - a support bot
   that makes up policy is worse than one that admits it does not know. */

function llm_answer(question, count)
{
    int best;
    char *fact;

    if (g_llm_open == 0)
        return "no model is open";

    best = llm_kb_best(question);
    if (best == 0)
        return "I have no information on that yet.";

    fact = llm_kb_fact(icpp_file_line(llm_kb_path(), best));
    if (count < 1)
        return fact;
    return isprint(fact, " ", llm_generate(fact, count));
}
