/* spotlight.c -- runs a Python program on the Spotlight VM model: ./spotlight [-v] < program.py
 * Run it in the model folder; it loads intelligence/ and every package in memory/, and -v also
 * prints each generated token. Built with -shared -DSPOTLIGHT_LIBRARY, spotlight.py loads it. */
#define _POSIX_C_SOURCE 200809L
#include <dirent.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef float real_t;
static void die(const char *what) { fprintf(stderr, "spotlight: %s\n", what); exit(1); }
static void *alloc(size_t n) { void *p = calloc(n ? n : 1, 1); if (!p) die("out of memory"); return p; }
static char folder[4096] = ".";
static const char *file(const char *name) { static char path[4200]; snprintf(path, sizeof path, "%s/%s", folder, name); return path; }

static char *slurp(const char *path, size_t *size) {
    FILE *f = path ? fopen(path, "rb") : stdin;
    if (!f) die(path);
    size_t cap = 1 << 16, n = 0, got; char *buf = alloc(cap);
    while ((got = fread(buf + n, 1, cap - n - 1, f)) > 0) if ((n += got) == cap - 1) buf = realloc(buf, cap *= 2);
    if (path) fclose(f);
    buf[n] = 0; if (size) *size = n;
    return buf;
}

/* ── config.json and safetensors ─────────────────────────────────────── */
static char *config;
static double number(const char *key) {
    char pattern[64]; snprintf(pattern, sizeof pattern, "\"%s\":", key);
    const char *p = strstr(config, pattern);
    if (!p) die(key);
    return strtod(p + strlen(pattern), NULL);
}

typedef struct { char dtype[4]; const char *data; int64_t count; } Tensor;
static int tensor(const char *file_bytes, const char *name, Tensor *out) {
    char pattern[96]; snprintf(pattern, sizeof pattern, "\"%s\":", name);
    uint64_t header; memcpy(&header, file_bytes, 8);
    const char *p = strstr(file_bytes + 8, pattern);
    if (!p || p > file_bytes + 8 + header) return 0;
    const char *dtype = strstr(p, "\"dtype\":\"") + 9; char *offsets = strstr(p, "\"data_offsets\":[") + 16;
    memset(out->dtype, 0, 4); memcpy(out->dtype, dtype, strchr(dtype, '"') - dtype);
    int64_t begin = strtoll(offsets, &offsets, 10), end = strtoll(offsets + 1, NULL, 10);
    int width = out->dtype[1] == '8' ? 1 : out->dtype[1] == '1' ? 2 : out->dtype[1] == '3' ? 4 : 8;
    out->data = file_bytes + 8 + header + begin; out->count = (end - begin) / width;
    return 1;
}
static double item(const Tensor *t, int64_t i) {
    double f; int32_t i32; int16_t i16;
    switch (t->dtype[0] == 'F' ? 'F' : t->dtype[1]) {
        case 'F': memcpy(&f, t->data + 8 * i, 8); return f;
        case '3': memcpy(&i32, t->data + 4 * i, 4); return i32;
        case '1': memcpy(&i16, t->data + 2 * i, 2); return i16;
        default: return t->dtype[0] == 'U' ? (double)(uint8_t)t->data[i] : (double)(int8_t)t->data[i];
    }
}
static real_t *load(const char *file_bytes, const char *name, int64_t count) {
    Tensor x; if (!tensor(file_bytes, name, &x) || x.count != count) die(name);
    real_t *out = alloc((size_t)count * sizeof(real_t));
    for (int64_t i = 0; i < count; i++) out[i] = (real_t)item(&x, i);
    return out;
}

/* ── the memory: one hash map from (table, x, y) to a cell's state ──── */
typedef struct { int64_t x, y; int32_t table, used; real_t s[2]; } Cell;
static Cell *cells; static size_t cells_cap, cells_n;
static const real_t ZERO[2] = {0, 0};

static size_t cell_slot(int table, int64_t x, int64_t y) {
    uint64_t h = (uint64_t)x * 0x9E3779B97F4A7C15ull ^ (uint64_t)y * 0xC2B2AE3D27D4EB4Full ^ (uint64_t)table * 0x165667B19E3779F9ull;
    size_t i = (size_t)(h ^ h >> 29) & (cells_cap - 1);
    while (cells[i].used && (cells[i].x != x || cells[i].y != y || cells[i].table != table)) i = (i + 1) & (cells_cap - 1);
    return i;
}
static real_t *cell_write(int table, int64_t x, int64_t y) {
    if (2 * (cells_n + 1) > cells_cap) {
        Cell *old = cells; size_t n = cells_cap;
        cells_cap = cells_cap ? 2 * cells_cap : 1 << 22; cells = alloc(cells_cap * sizeof(Cell));
        for (size_t i = 0; i < n; i++) if (old[i].used) cells[cell_slot(old[i].table, old[i].x, old[i].y)] = old[i];
        free(old);
    }
    size_t i = cell_slot(table, x, y);
    if (!cells[i].used) { cells[i].x = x; cells[i].y = y; cells[i].table = table; cells[i].used = 1; cells_n++; }
    return cells[i].s;
}
static const real_t *cell_read(int table, int64_t x, int64_t y) {
    size_t i = cell_slot(table, x, y);
    return cells[i].used ? cells[i].s : ZERO;
}

/* ── the model ───────────────────────────────────────────────────────── */
/* The Spotlight kernel: phi_z(x) = b(x - z) / sqrt(sum_j b(x - j)^2) over the three nearest cells,
 * with b(t) = cos^2(pi t / 2) for |t| < 1. */
static void phi(real_t x, int64_t *z0, real_t w[3]) {
    int64_t c = (int64_t)floor(x + 0.5); real_t norm = 0;
    for (int i = 0; i < 3; i++) {
        real_t t = x - (real_t)(c + i - 1), b = fabsf(t) < 1 ? cosf(3.14159265f * t / 2) : 0;
        w[i] = b * b; norm += w[i] * w[i];
    }
    for (int i = 0; i < 3; i++) w[i] /= sqrtf(norm);
    *z0 = c - 1;
}
static real_t sigmoid(real_t z) { return 1 / (1 + expf(-z)); }
static real_t dot(const real_t *row, const real_t *x, int cols) {   /* 8 lanes, so it vectorizes */
    real_t a[8] = {0}, acc; int c = 0;
    for (; c + 8 <= cols; c += 8) for (int k = 0; k < 8; k++) a[k] += row[c + k] * x[c + k];
    acc = ((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]));
    for (; c < cols; c++) acc += row[c] * x[c];
    return acc;
}
static void matvec(const real_t *m, int rows, int cols, const real_t *x, const real_t *bias, real_t *y) {
    for (int r = 0; r < rows; r++) y[r] = bias ? dot(m + (size_t)r * cols, x, cols) + bias[r] : dot(m + (size_t)r * cols, x, cols);
}

typedef struct { real_t *qkv_w, *qkv_b, *out_w, *ff_in_w, *ff_out_w; } Layer;
static struct {
    int D, H, layers, vocab, proj_out, stop_id, n_out, d_ffn, rank;
    Layer *layer; real_t *embed, *basis, *coeff; int32_t *out_ids; char **tokens;
} M;

/* One head, one token: a delta-rule write into the 3x3 cells around the key, then a read of the
 * 3x3 cells around the query. */
static void memory_step(int table, const real_t *p, int h, real_t *out) {
    const int H = M.H; const real_t *v = p + 6 * H + 2 * h;
    real_t beta = sigmoid(p[4 * H + h]), alpha = sigmoid(-p[5 * H + h]), wx[3], wy[3];
    int64_t x0, y0;
    phi(p[2 * h], &x0, wx); phi(p[2 * h + 1], &y0, wy);
    for (int i = 0; i < 3; i++) for (int j = 0; j < 3; j++) {
        real_t w = wx[i] * wy[j] * beta;
        if (w == 0.0) continue;
        real_t *S = cell_write(table, x0 + i, y0 + j);
        real_t d0 = (v[0] - alpha * S[0]) * w, d1 = (v[1] - alpha * S[1]) * w;
        S[0] += d0; S[1] += d1;
    }
    phi(p[2 * H + 2 * h], &x0, wx); phi(p[2 * H + 2 * h + 1], &y0, wy);
    out[0] = out[1] = 0.0;
    for (int i = 0; i < 3; i++) for (int j = 0; j < 3; j++) {
        const real_t *S = cell_read(table, x0 + i, y0 + j);
        out[0] += wx[i] * wy[j] * S[0]; out[1] += wx[i] * wy[j] * S[1];
    }
}

/* One token through every layer, then the most likely next token. */
static int32_t forward(int32_t token) {
    static real_t x[64], proj[512], mem[64], res[64], pre[256], gate[128], feat[16];
    const int D = M.D;
    memcpy(x, M.embed + (size_t)token * D, D * sizeof(real_t));
    for (int l = 0; l < M.layers; l++) {
        const Layer *L = &M.layer[l];
        matvec(L->qkv_w, M.proj_out, D, x, L->qkv_b, proj);
        for (int h = 0; h < M.H; h++) memory_step(l * M.H + h, proj, h, mem + 2 * h);
        matvec(L->out_w, D, 2 * M.H, mem, NULL, res);
        for (int i = 0; i < D; i++) x[i] += res[i];
        matvec(L->ff_in_w, 2 * M.d_ffn, D, x, NULL, pre);
        for (int j = 0; j < M.d_ffn; j++) gate[j] = (pre[j] > 0.0 ? pre[j] : 0.0) * pre[M.d_ffn + j];
        matvec(L->ff_out_w, D, M.d_ffn, gate, NULL, res);
        for (int i = 0; i < D; i++) x[i] += res[i];
    }
    matvec(M.basis, M.rank, D, x, NULL, feat);
    real_t best = -HUGE_VAL; int32_t best_row = 0;
    for (int r = 0; r < M.n_out; r++) {
        real_t acc = 0.0;
        for (int k = 0; k < M.rank; k++) acc += M.coeff[(size_t)r * M.rank + k] * feat[k];
        if (acc > best) { best = acc; best_row = r; }
    }
    return M.out_ids[best_row];
}

/* ── loading ─────────────────────────────────────────────────────────── */
static void poke(int64_t address, const uint8_t *bytes, size_t n) {   /* bytes into the memory's cells */
    int table = number("memory_layer") * M.H + number("memory_head"); double spacing = number("spacing");
    for (size_t i = 0; i < n; i++, address++) {
        real_t *s = cell_write(table, (int64_t)(spacing * (double)(address % 65536)), (int64_t)(spacing * (double)(address / 65536)));
        s[0] = (real_t)bytes[i]; s[1] = 0.0;
    }
}
static void poke_word(int64_t address, uint32_t v) { uint8_t b[4] = {v & 255, v >> 8 & 255, v >> 16 & 255, v >> 24}; poke(address, b, 4); }

int spotlight_load(const char *dir) {
    snprintf(folder, sizeof folder, "%s", dir);
    config = slurp(file("intelligence/config.json"), NULL);
    M.D = number("d_model"); M.H = number("n_heads"); M.layers = number("n_layers"); M.vocab = number("vocab_size");
    M.proj_out = number("proj_out"); M.stop_id = number("stop_id"); M.n_out = number("n_out");
    M.d_ffn = number("d_ffn"); M.rank = number("head_rank");
    if (number("d_rest") != 0 || number("d_v") != 2 || number("alpha_rows") != 1) die("unsupported model shape");
    char *vocab = slurp(file("intelligence/vocab.json"), NULL), *p = vocab, name[64];
    M.tokens = alloc(M.vocab * sizeof(char *));
    for (int i = 0; i < M.vocab; i++) {
        p = strchr(p, '"') + 1; char *q = M.tokens[i] = alloc(64);
        while (*p != '"') { if (*p == '\\') p++; *q++ = *p++; }
        p++;
    }
    char *w = slurp(file("intelligence/model.safetensors"), NULL);
    M.layer = alloc(M.layers * sizeof(Layer));
    for (int l = 0; l < M.layers; l++) {
        Layer *L = &M.layer[l];
        snprintf(name, sizeof name, "layers.%d.qkv_w", l); L->qkv_w = load(w, name, (int64_t)M.proj_out * M.D);
        snprintf(name, sizeof name, "layers.%d.qkv_b", l); L->qkv_b = load(w, name, M.proj_out);
        snprintf(name, sizeof name, "layers.%d.out_w", l); L->out_w = load(w, name, (int64_t)M.D * 2 * M.H);
        snprintf(name, sizeof name, "layers.%d.ff_in_w", l); L->ff_in_w = load(w, name, (int64_t)2 * M.d_ffn * M.D);
        snprintf(name, sizeof name, "layers.%d.ff_out_w", l); L->ff_out_w = load(w, name, (int64_t)M.D * M.d_ffn);
    }
    M.embed = load(w, "embed_tokens.weight", (int64_t)M.vocab * M.D);
    M.basis = load(w, "head_basis", (int64_t)M.rank * M.D);
    M.coeff = load(w, "head_coeff", (int64_t)M.n_out * M.rank);
    real_t *ids = load(w, "out_ids", M.n_out);
    M.out_ids = alloc(M.n_out * sizeof(int32_t));
    for (int i = 0; i < M.n_out; i++) M.out_ids[i] = (int32_t)ids[i];
    free(vocab); free(w); free(ids);
    return M.vocab;
}

/* A fresh memory: MicroPython's code, tables and boot state, then the packages in one registry. */
static void load_memory(char **packages, int n) {
    free(cells); cells = NULL; cells_cap = cells_n = 0;
    char *core = slurp(file("memory/core.safetensors"), NULL), name[64];
    for (int t = 0; t < M.layers * M.H; t++) {
        Tensor xy, value;
        snprintf(name, sizeof name, "cells.%d.%d.xy", t / M.H, t % M.H); if (!tensor(core, name, &xy)) continue;
        snprintf(name, sizeof name, "cells.%d.%d.value", t / M.H, t % M.H); if (!tensor(core, name, &value)) continue;
        for (int64_t i = 0; i < xy.count / 2; i++) {
            real_t *s = cell_write(t, (int64_t)item(&xy, 2 * i), (int64_t)item(&xy, 2 * i + 1));
            s[0] = (real_t)item(&value, 2 * i); s[1] = (real_t)item(&value, 2 * i + 1);
        }
    }
    free(core);
    if (n == 0) return;
    int64_t base = number("registry_base"), size = 16, pages, modules = 0;
    for (int i = 0; i < n; i++) {
        char *bytes = slurp(packages[i], NULL); Tensor b;
        if (!tensor(bytes, "bundle", &b) || memcmp(b.data, "SPKG", 4)) die("not a package");
        poke(base + size, (const uint8_t *)b.data + 16, b.count - 16);
        uint32_t count; memcpy(&count, b.data + 12, 4);
        size += b.count - 16; modules += count;
        free(bytes);
    }
    uint8_t header[8] = {'S', 'P', 'K', 'G', 1};
    poke(base, header, 8); poke_word(base + 8, (uint32_t)size); poke_word(base + 12, (uint32_t)modules);
    poke_word(number("registry_root_word"), (uint32_t)base);
    pages = (base + size + 65535) / 65536;
    poke_word(number("memory_pages_word"), (uint32_t)(pages > number("image_pages") ? pages : number("image_pages")));
}

/* ── running ─────────────────────────────────────────────────────────── */
static int32_t *kept; static int64_t kept_n, kept_cap; static int printing = 0, verbose = 0;

/* Runs a prompt from a fresh memory until the stop token (or max_new tokens). The command line
 * prints the output as it comes; the library keeps the tokens that carry it, and the stop token. */
int64_t spotlight_generate(const int32_t *prompt, int64_t n, int64_t max_new, char **packages, int n_packages) {
    load_memory(packages, n_packages);
    int32_t next = 0; kept_n = 0;
    for (int64_t i = 0; i < n; i++) next = forward(prompt[i]);
    for (int64_t count = 1;; count++, next = forward(next)) {
        const char *t = M.tokens[next]; int out = !strncmp(t, "out(", 4);
        if (verbose) fprintf(stderr, "%s ", t);
        if (printing && out) { putchar(strlen(t) == 6 ? t[4] : (int)strtol(t + 4, NULL, 16)); fflush(stdout); }
        if (!printing && (out || next == M.stop_id)) {
            if (kept_n == kept_cap) kept = realloc(kept, (kept_cap = kept_cap ? 2 * kept_cap : 1024) * sizeof(int32_t));
            kept[kept_n++] = next;
        }
        if (next == M.stop_id || count == max_new) return count;
    }
}
const int32_t *spotlight_kept(int64_t *n) { *n = kept_n; return kept; }

#ifndef SPOTLIGHT_LIBRARY
static int32_t token(const char *name) {
    for (int32_t i = 0; i < M.vocab; i++) if (!strcmp(M.tokens[i], name)) return i;
    die(name); return -1;
}
static int is_package(const struct dirent *e) {
    size_t n = strlen(e->d_name);
    return n > 12 && !strcmp(e->d_name + n - 12, ".safetensors") && strcmp(e->d_name, "core.safetensors");
}
int main(int argc, char **argv) {
    verbose = argc > 1 && !strcmp(argv[1], "-v"); printing = 1;
    spotlight_load(".");
    struct dirent **list; int n_packages = scandir("memory", &list, is_package, alphasort);
    char **packages = alloc((n_packages > 0 ? n_packages : 1) * sizeof(char *));
    for (int i = 0; i < n_packages; i++) { packages[i] = alloc(300); snprintf(packages[i], 300, "memory/%s", list[i]->d_name); }
    /* The prompt: where the source goes, its bytes (and a NUL) page by page, then execute. */
    size_t n; char *src = slurp(NULL, &n), hex[16];
    int64_t input = number("input_base"), pages = (input % 65536 + (int64_t)n) / 65536 + 1, count = 0;
    int32_t *prompt = alloc((n + 16 * pages + 8) * sizeof(int32_t));
    for (int64_t r = pages - 1; r >= 0; r--) {
        uint32_t address = (uint32_t)(r ? (input / 65536 + r) * 65536 : input);
        for (int lane = 0; lane < 4; lane++) { snprintf(hex, sizeof hex, "%02x", address >> 8 * lane & 255); prompt[count++] = token(hex); }
        prompt[count++] = token("page");
    }
    for (int64_t i = 0; i <= (int64_t)n; i++) {
        unsigned char b = i < (int64_t)n ? src[i] : 0;
        if (b >= 128) die("the program must be ASCII");
        if (i && (input + i) % 65536 == 0) prompt[count++] = token("unpage");
        if (b > 0x20 && b < 0x7f) { hex[0] = b; hex[1] = 0; } else snprintf(hex, sizeof hex, "ascii(%02x)", b);
        prompt[count++] = token(hex);
    }
    prompt[count++] = token("execute");
    count = spotlight_generate(prompt, count, INT64_MAX, packages, n_packages > 0 ? n_packages : 0);
    fprintf(stderr, "%s%lld tokens generated\n", verbose ? "\n" : "", (long long)count);
    return 0;
}
#endif