/* test_vindex.c — build + RUN gate for the M8 HNSW vector index.
 *
 * Covers: recall@10 vs brute-force oracle, brute-force-vs-index speedup,
 * correctness edge cases (k>N, identical vectors, self-query, zero vector),
 * determinism (seeded PRNG → identical graphs), and vindex_build_from_store
 * over a real engram_store on-disk file.
 *
 * Pure C11; links engram_vindex.c + engram_store.c; -lm. ASan/UBSan clean.
 */
#include "engram_vindex.h"
#include "engram_store.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <time.h>
#include <unistd.h>

#define DIM 768

static int g_fail = 0;
/* VINDEX_QUICK=1 shrinks the two large builds so the ASan/UBSan pass (which runs
 * ~5-10x slower) stays fast — memory-safety is size-independent. The perf numbers
 * (recall gate + speedup) come from the un-sanitized, full-size pass. */
static int g_quick = 0;
static int envint(const char* k, int dflt){ const char* s=getenv(k); return s?atoi(s):dflt; }
#define CHECK(cond, msg) do{ if(!(cond)){ printf("  FAIL: %s\n", msg); g_fail=1; } else { printf("  ok: %s\n", msg); } }while(0)

/* deterministic test PRNG (splitmix64) */
static uint64_t rng_state = 0xABCDEF0123456789ULL;
static uint64_t xrng(uint64_t* s){
    uint64_t z=(*s+=0x9E3779B97F4A7C15ULL);
    z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL;
    return z^(z>>31);
}
static float frand(uint64_t* s){ return (float)((xrng(s)>>11)*(1.0/9007199254740992.0)) - 0.5f; }

static double now_s(void){
    struct timespec t; clock_gettime(CLOCK_MONOTONIC,&t);
    return t.tv_sec + t.tv_nsec*1e-9;
}

/* fill vec[N*DIM]: mostly random, some clustered groups (center + small noise). */
static void gen_vectors(float* v, int N, uint64_t seed){
    uint64_t s = seed;
    int clustered = N/5;                 /* last fifth is clustered */
    int ncenters = 20;
    float* centers = (float*)malloc((size_t)ncenters*DIM*sizeof(float));
    for (int c=0;c<ncenters;c++) for(int d=0;d<DIM;d++) centers[c*DIM+d]=frand(&s);
    for (int i=0;i<N;i++){
        if (i < N-clustered){
            for (int d=0;d<DIM;d++) v[i*DIM+d]=frand(&s);
        } else {
            int c = (int)(xrng(&s)%ncenters);
            for (int d=0;d<DIM;d++) v[i*DIM+d]=centers[c*DIM+d] + 0.05f*frand(&s);
        }
    }
    free(centers);
}

static float cosdist(const float* a, const float* b){
    double da=0,db=0,dot=0;
    for(int i=0;i<DIM;i++){ da+=(double)a[i]*a[i]; db+=(double)b[i]*b[i]; dot+=(double)a[i]*b[i]; }
    if (da<=0||db<=0) return 1.0f;
    return (float)(1.0 - dot/(sqrt(da)*sqrt(db)));
}

/* brute-force top-k node ids into ids[k] (ascending distance). */
static void brute_topk(const float* v, int N, const float* q, int k, int* ids){
    float* bd = (float*)malloc((size_t)k*sizeof(float));
    for (int i=0;i<k;i++){ ids[i]=-1; bd[i]=1e30f; }
    for (int i=0;i<N;i++){
        float d = cosdist(q, v+(size_t)i*DIM);
        if (d < bd[k-1]){
            int p=k-1;
            while (p>0 && bd[p-1]>d){ bd[p]=bd[p-1]; ids[p]=ids[p-1]; p--; }
            bd[p]=d; ids[p]=i;
        }
    }
    free(bd);
}

/* ── Test 1: recall@10 vs brute force + latency/recall tradeoff ────────────── */
static void test_recall(void){
    int N=envint("VINDEX_N_RECALL", g_quick?1500:5000), Q=200, K=10;
    printf("\n== Test 1: recall@10 vs brute force (N=%d, DIM=768) ==\n", N);
    float* v = (float*)malloc((size_t)N*DIM*sizeof(float));
    gen_vectors(v, N, 111);

    double t0=now_s();
    VIndex* ix = vindex_create(DIM, VINDEX_DEFAULT_M, VINDEX_DEFAULT_EF_CONSTRUCTION);
    for (int i=0;i<N;i++) vindex_insert(ix, (uint64_t)i, v+(size_t)i*DIM);
    double build_s = now_s()-t0;
    printf("  build: %d vectors in %.2fs (M=%d, ef_construction=%d)\n",
           N, build_s, VINDEX_DEFAULT_M, VINDEX_DEFAULT_EF_CONSTRUCTION);

    /* queries: half random, half near a real vector (perturbed). */
    float* qs = (float*)malloc((size_t)Q*DIM*sizeof(float));
    uint64_t s=999;
    for (int i=0;i<Q;i++){
        if (i<Q/2) for(int d=0;d<DIM;d++) qs[i*DIM+d]=frand(&s);
        else { int base=(int)(xrng(&s)%N); for(int d=0;d<DIM;d++) qs[i*DIM+d]=v[base*DIM+d]+0.03f*frand(&s); }
    }

    /* oracle */
    int* oracle = (int*)malloc((size_t)Q*K*sizeof(int));
    for (int i=0;i<Q;i++) brute_topk(v, N, qs+(size_t)i*DIM, K, oracle+(size_t)i*K);

    int efs[] = { 10, 32, 64, 128 };
    for (int e=0;e<4;e++){
        int ef=efs[e];
        uint64_t ids[64]; float dd[64];
        int hits=0;
        double qt0=now_s();
        for (int i=0;i<Q;i++){
            int n=vindex_search(ix, qs+(size_t)i*DIM, K, ef, ids, dd);
            for (int a=0;a<n;a++) for(int b=0;b<K;b++) if((int)ids[a]==oracle[i*K+b]){ hits++; break; }
        }
        double qs_ms = (now_s()-qt0)*1000.0/Q;
        double recall = (double)hits/(Q*K);
        printf("  ef_search=%-4d recall@10=%.4f  latency=%.3f ms/query\n", ef, recall, qs_ms);
        if (ef==VINDEX_DEFAULT_EF_SEARCH && !g_quick)
            CHECK(recall >= 0.90, "recall@10 >= 0.90 at default ef_search=128");
    }
    free(oracle); free(qs); free(v); vindex_free(ix);
}

/* ── Test 2: speedup vs brute force ───────────────────────────────────────── */
static void speedup_at(int N){
    int Q=100, K=10;
    float* v=(float*)malloc((size_t)N*DIM*sizeof(float));
    gen_vectors(v,N,222);
    VIndex* ix=vindex_create(DIM,16,200);
    double bt0=now_s();
    for(int i=0;i<N;i++) vindex_insert(ix,(uint64_t)i,v+(size_t)i*DIM);
    printf("  N=%d build=%.2fs\n", N, now_s()-bt0);

    float* qs=(float*)malloc((size_t)Q*DIM*sizeof(float));
    uint64_t s=333; for(int i=0;i<Q*DIM;i++) qs[i]=frand(&s);

    /* brute force */
    int scratch[16];
    double b0=now_s();
    for(int i=0;i<Q;i++) brute_topk(v,N,qs+(size_t)i*DIM,K,scratch);
    double bf=(now_s()-b0)/Q;

    /* index */
    uint64_t ids[16]; float dd[16];
    double i0=now_s();
    for(int i=0;i<Q;i++) vindex_search(ix,qs+(size_t)i*DIM,K,64,ids,dd);
    double iq=(now_s()-i0)/Q;

    printf("  N=%d  brute=%.4f ms/q  index=%.4f ms/q  speedup=%.1fx\n",
           N, bf*1000, iq*1000, bf/iq);
    CHECK(iq < bf, "index query faster than brute force");
    free(qs); free(v); vindex_free(ix);
}
static void test_speedup(void){
    printf("\n== Test 2: brute-force vs index speedup ==\n");
    speedup_at(g_quick?2000:5000);
    speedup_at(envint("VINDEX_N_BIG", g_quick?3000:20000));
}

/* ── Test 3: edge cases ───────────────────────────────────────────────────── */
static void test_edges(void){
    printf("\n== Test 3: correctness edge cases ==\n");
    /* k larger than node count */
    {
        VIndex* ix=vindex_create(DIM,16,200);
        float vec[DIM]; uint64_t s=1;
        for(int i=0;i<3;i++){ for(int d=0;d<DIM;d++) vec[d]=frand(&s); vindex_insert(ix,(uint64_t)i,vec); }
        uint64_t ids[50]; float dd[50];
        int n=vindex_search(ix, vec, 50, 64, ids, dd);
        CHECK(n==3, "k > node count returns exactly node-count results");
        vindex_free(ix);
    }
    /* duplicate / identical vectors */
    {
        VIndex* ix=vindex_create(DIM,16,200);
        float a[DIM]; uint64_t s=2; for(int d=0;d<DIM;d++) a[d]=frand(&s);
        for(int i=0;i<10;i++) vindex_insert(ix,(uint64_t)i,a);   /* all identical */
        float b[DIM]; for(int d=0;d<DIM;d++) b[d]=frand(&s);
        vindex_insert(ix,100,b);
        uint64_t ids[5]; float dd[5];
        int n=vindex_search(ix,a,5,64,ids,dd);
        CHECK(n==5, "identical-vector index returns k results");
        CHECK(dd[0] < 1e-4f, "top-1 distance ~0 for a duplicated vector");
        vindex_free(ix);
    }
    /* query equal to an indexed vector returns itself as top-1, dist ~0 */
    {
        VIndex* ix=vindex_create(DIM,16,200);
        int N=500; float* v=(float*)malloc((size_t)N*DIM*sizeof(float)); gen_vectors(v,N,7);
        for(int i=0;i<N;i++) vindex_insert(ix,(uint64_t)(1000+i),v+(size_t)i*DIM);
        int probe=137;
        uint64_t ids[3]; float dd[3];
        int n=vindex_search(ix, v+(size_t)probe*DIM, 3, 64, ids, dd);
        CHECK(n>=1 && ids[0]==(uint64_t)(1000+probe), "self-query returns itself as top-1");
        CHECK(dd[0] < 1e-4f, "self-query top-1 distance ~0");
        free(v); vindex_free(ix);
    }
    /* zero vector: no NaN, handled */
    {
        VIndex* ix=vindex_create(DIM,16,200);
        float z[DIM]; memset(z,0,sizeof z);
        float a[DIM]; uint64_t s=3; for(int d=0;d<DIM;d++) a[d]=frand(&s);
        vindex_insert(ix,0,z); vindex_insert(ix,1,a);
        uint64_t ids[2]; float dd[2];
        int n=vindex_search(ix, z, 2, 64, ids, dd);            /* zero query */
        int nan=0; for(int i=0;i<n;i++) if(isnan(dd[i])||isinf(dd[i])) nan=1;
        CHECK(n>=1 && !nan, "zero vector query produces no NaN/Inf");
        n=vindex_search(ix, a, 2, 64, ids, dd);                /* zero indexed */
        nan=0; for(int i=0;i<n;i++) if(isnan(dd[i])||isinf(dd[i])) nan=1;
        CHECK(!nan, "indexed zero vector produces no NaN/Inf");
        vindex_free(ix);
    }
}

/* ── Test 4: determinism ──────────────────────────────────────────────────── */
static void test_determinism(void){
    printf("\n== Test 4: determinism (seeded PRNG → identical results) ==\n");
    int N=1500;
    float* v=(float*)malloc((size_t)N*DIM*sizeof(float)); gen_vectors(v,N,55);
    uint64_t ids1[10],ids2[10]; float d1[10],d2[10];
    int identical=1;
    for (int build=0; build<2; build++){
        VIndex* ix=vindex_create(DIM,16,200);
        for(int i=0;i<N;i++) vindex_insert(ix,(uint64_t)i,v+(size_t)i*DIM);
        /* probe several queries */
        for (int q=0;q<20;q++){
            uint64_t* ida = build? ids2 : ids1; float* da = build? d2 : d1;
            vindex_search(ix, v+(size_t)(q*37%N)*DIM, 10, 64, ida, da);
            if (build==1){
                /* re-run build-0 query stored? simpler: compare within-run below */
            }
        }
        vindex_free(ix);
    }
    /* Proper comparison: run two fresh builds, same single query. */
    identical=1;
    for (int q=0;q<25;q++){
        int qi=(q*61)%N;
        VIndex* a=vindex_create(DIM,16,200); for(int i=0;i<N;i++) vindex_insert(a,(uint64_t)i,v+(size_t)i*DIM);
        VIndex* b=vindex_create(DIM,16,200); for(int i=0;i<N;i++) vindex_insert(b,(uint64_t)i,v+(size_t)i*DIM);
        int na=vindex_search(a, v+(size_t)qi*DIM,10,64,ids1,d1);
        int nb=vindex_search(b, v+(size_t)qi*DIM,10,64,ids2,d2);
        if (na!=nb) identical=0;
        for(int i=0;i<na;i++) if(ids1[i]!=ids2[i] || d1[i]!=d2[i]) identical=0;
        vindex_free(a); vindex_free(b);
    }
    CHECK(identical, "two independent builds give byte-identical query results");
    free(v);
}

/* ── Test 5: build_from_store ─────────────────────────────────────────────── */
static void test_build_from_store(void){
    printf("\n== Test 5: vindex_build_from_store over a real engram_store ==\n");
    char path[256];
    snprintf(path,sizeof path,"/tmp/vindex_test_store_%d.engram",(int)getpid());
    unlink(path);
    EngramPagedStore* st = store_create(path);
    if (!st){ printf("  FAIL: store_create\n"); g_fail=1; return; }

    int N=300;
    float* v=(float*)malloc((size_t)N*DIM*sizeof(float)); gen_vectors(v,N,88);
    for (int i=0;i<N;i++){
        StoreNode n; memset(&n,0,sizeof n);
        char id[32]; snprintf(id,sizeof id,"node-%d",i);
        n.id=id; n.content="x"; n.node_type="concept"; n.tier="Semantic";
        n.emb = v+(size_t)i*DIM; n.emb_dim=DIM;
        if (store_put_node(st,&n)!=0){ printf("  FAIL: put_node %d\n",i); g_fail=1; }
    }
    /* a node WITHOUT an emb — must be skipped by build_from_store. */
    { StoreNode n; memset(&n,0,sizeof n); n.id=(char*)"no-emb"; n.content="y"; n.node_type="concept"; n.tier="Semantic";
      store_put_node(st,&n); }
    store_close(st);

    VIndex* ix = vindex_create(DIM,16,200);
    char** ids=NULL; int nids=0;
    int ins = vindex_build_from_store(ix, path, &ids, &nids);
    printf("  build_from_store inserted %d vectors (expected %d; 1 emb-less skipped)\n", ins, N);
    CHECK(ins==N, "build_from_store inserts exactly the emb'd nodes");
    CHECK((size_t)ins==vindex_size(ix), "index size matches insert count");

    /* query with a known vector → must return its own node id as top-1. */
    int probe=42;
    uint64_t rids[5]; float dd[5];
    int n=vindex_search(ix, v+(size_t)probe*DIM, 5, 64, rids, dd);
    int correct = (n>=1 && rids[0]<(uint64_t)nids && strcmp(ids[rids[0]], "node-42")==0);
    printf("  query for node-42's vector → top-1 id=%s dist=%.5f\n",
           (n>=1 && rids[0]<(uint64_t)nids)? ids[rids[0]] : "?", n?dd[0]:-1);
    CHECK(correct, "build_from_store query resolves to the right node id");
    CHECK(n>=1 && dd[0]<1e-4f, "top-1 distance ~0 for exact stored vector");

    for (int i=0;i<nids;i++) free(ids[i]);
    free(ids); free(v); vindex_free(ix); unlink(path);
}

int main(void){
    (void)rng_state;
    g_quick = envint("VINDEX_QUICK", 0);
    printf("=== engram_vindex (HNSW) test suite ===%s\n", g_quick?" [QUICK]":"");
    test_recall();
    test_speedup();
    test_edges();
    test_determinism();
    test_build_from_store();
    printf("\n=== %s ===\n", g_fail? "FAILURES PRESENT" : "ALL TESTS PASSED");
    return g_fail;
}