/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_rosa_db_qa.c — Phase 4: Model answering with ROSA context
 *
 * The KEY TEST: Can the 3.6B model answer questions about past conversations
 * when given ROSA-retrieved context?
 *
 * For each question, runs:
 *   A) Baseline: just the question (no context)
 *   B) ROSA: retrieved context + question
 *
 * Usage:
 *   ./examples/test_rosa_db_qa \
 *     --model <path> --vocab <path> --db <path> \
 *     --context-window 256 --max-tokens 64 --quantize
 */

#include "rosa/rosa.h"
#include "rwkvx/rwkvx.h"
#include "data/data.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define MAX_TOKENS_PER_CONV (1024 * 64)
#define MAX_QUERY_TOKENS 512
#define MAX_PROMPT_TOKENS (1024 * 4)

/* ---- Index builder (reused from Phase 3) ---- */

typedef struct {
    sa_token_automaton *sa;
    sa_position_map *pmap;
    int total_tokens;
    int n_conv;
    int *conv_start;
    int *conv_len;
    uint32_t *all_tokens;
    int all_tokens_cap;
} rosa_index;

static rosa_index* build_index(fd_dataset *ds, rx_tokenizer *tok) {
    rosa_index *idx = (rosa_index *)calloc(1, sizeof(rosa_index));
    idx->n_conv = fd_len(ds);
    idx->sa = sa_token_new(1024, 65536);
    idx->pmap = sa_position_map_new(4096);
    idx->conv_start = (int *)malloc(idx->n_conv * sizeof(int));
    idx->conv_len = (int *)malloc(idx->n_conv * sizeof(int));
    idx->all_tokens_cap = 1024 * 1024;
    idx->all_tokens = (uint32_t *)malloc(idx->all_tokens_cap * sizeof(uint32_t));

    uint32_t *buf = (uint32_t *)malloc(MAX_TOKENS_PER_CONV * sizeof(uint32_t));

    for (int c = 0; c < idx->n_conv; c++) {
        fd_sample *sample = fd_get(ds, c);
        const char *text = fd_sample_get(sample, "text");
        if (!text) { idx->conv_start[c] = idx->total_tokens; idx->conv_len[c] = 0; continue; }

        int n = rx_tokenizer_encode(tok, text, buf, MAX_TOKENS_PER_CONV);
        if (n <= 0) { idx->conv_start[c] = idx->total_tokens; idx->conv_len[c] = 0; continue; }

        idx->conv_start[c] = idx->total_tokens;
        idx->conv_len[c] = n;

        while (idx->total_tokens + n > idx->all_tokens_cap) {
            idx->all_tokens_cap *= 2;
            idx->all_tokens = (uint32_t *)realloc(idx->all_tokens,
                idx->all_tokens_cap * sizeof(uint32_t));
        }

        for (int t = 0; t < n; t++) {
            int gpos = idx->total_tokens + t;
            sa_token_extend(idx->sa, (int)buf[t], gpos);
            sa_position_map_append(idx->pmap, c, 0, t);
            idx->all_tokens[gpos] = buf[t];
        }
        idx->total_tokens += n;
    }

    free(buf);
    return idx;
}

static void free_index(rosa_index *idx) {
    sa_token_free(idx->sa);
    sa_position_map_free(idx->pmap);
    free(idx->conv_start);
    free(idx->conv_len);
    free(idx->all_tokens);
    free(idx);
}

/* ---- Retrieve context for a query ---- */

static char* retrieve_context(rosa_index *idx, rx_tokenizer *tok,
                               const char *query, int context_window) {
    uint32_t qtokens[MAX_QUERY_TOKENS];
    int qlen = rx_tokenizer_encode(tok, query, qtokens, MAX_QUERY_TOKENS);
    if (qlen <= 0) return NULL;

    int state = 0, match_len = 0;
    int best_endpos = -1, best_match = 0;

    for (int i = 0; i < qlen; i++) {
        int ep = sa_token_query_step(idx->sa, (int)qtokens[i], &state, &match_len);
        if (match_len > best_match) {
            best_match = match_len;
            best_endpos = ep;
        }
    }

    if (best_endpos < 0) return NULL;

    const sa_position_info *info = sa_position_map_get(idx->pmap, best_endpos);
    if (!info) return NULL;

    int conv_start = idx->conv_start[info->conv_idx];
    int conv_len = idx->conv_len[info->conv_idx];

    int match_end = info->token_offset;
    int ctx_start = match_end - best_match - context_window / 2;
    if (ctx_start < 0) ctx_start = 0;
    int ctx_end = match_end + context_window / 2;
    if (ctx_end > conv_len) ctx_end = conv_len;

    int ctx_len = ctx_end - ctx_start;
    if (ctx_len <= 0) return NULL;

    return rx_tokenizer_decode(tok, &idx->all_tokens[conv_start + ctx_start], ctx_len);
}

/* ---- Generate answer from model ---- */

static char* generate_answer(rx_model *model, rx_tokenizer *tok,
                              const char *prompt, int max_tokens) {
    rx_state *state = rx_state_new(model);

    /* Tokenize and prefill prompt */
    uint32_t *prompt_tokens = (uint32_t *)malloc(MAX_PROMPT_TOKENS * sizeof(uint32_t));
    int prompt_len = rx_tokenizer_encode(tok, prompt, prompt_tokens, MAX_PROMPT_TOKENS);

    tm_tensor *logits = NULL;
    for (int i = 0; i < prompt_len; i++) {
        if (logits) tm_free(logits);
        logits = rx_forward(model, prompt_tokens[i], state);
    }

    /* Generate tokens */
    uint32_t *out_tokens = (uint32_t *)malloc(max_tokens * sizeof(uint32_t));
    int n_out = 0;

    rx_sampling_params params = rx_sampling_defaults();
    params.temperature = 0.0f;  /* greedy for reproducibility */
    params.top_p = 1.0f;

    rx_config cfg = rx_model_config(model);

    for (int i = 0; i < max_tokens; i++) {
        uint32_t tok_id = rx_sample(logits, &params, NULL, cfg.vocab_size);
        tm_free(logits);
        logits = NULL;

        if (tok_id == 0) break;  /* EOS */
        out_tokens[n_out++] = tok_id;

        logits = rx_forward(model, tok_id, state);
    }

    if (logits) tm_free(logits);

    char *answer = NULL;
    if (n_out > 0)
        answer = rx_tokenizer_decode(tok, out_tokens, n_out);
    else
        answer = strdup("(no output)");

    free(prompt_tokens);
    free(out_tokens);
    rx_state_free(state);
    return answer;
}

/* ---- Main ---- */

int main(int argc, char **argv) {
    const char *model_path = NULL;
    const char *vocab_path = NULL;
    const char *db_path = NULL;
    int context_window = 256;
    int max_tokens = 64;
    int quantize = 0;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--model") == 0 && i + 1 < argc)
            model_path = argv[++i];
        else if (strcmp(argv[i], "--vocab") == 0 && i + 1 < argc)
            vocab_path = argv[++i];
        else if (strcmp(argv[i], "--db") == 0 && i + 1 < argc)
            db_path = argv[++i];
        else if (strcmp(argv[i], "--context-window") == 0 && i + 1 < argc)
            context_window = atoi(argv[++i]);
        else if (strcmp(argv[i], "--max-tokens") == 0 && i + 1 < argc)
            max_tokens = atoi(argv[++i]);
        else if (strcmp(argv[i], "--quantize") == 0)
            quantize = 1;
    }

    if (!model_path || !vocab_path || !db_path) {
        fprintf(stderr, "Usage: %s --model <path> --vocab <path> --db <path> "
                "[--context-window N] [--max-tokens N] [--quantize]\n", argv[0]);
        return 1;
    }

    /* Load tokenizer */
    printf("Loading tokenizer...\n");
    rx_tokenizer *tok = rx_tokenizer_new(vocab_path);
    if (!tok) { fprintf(stderr, "ERROR: tokenizer\n"); return 1; }

    /* Load conversations and build index */
    printf("Loading conversations...\n");
    fd_dataset *ds = fd_load_chat_sqlite(db_path, 0);
    if (!ds) { fprintf(stderr, "ERROR: database\n"); return 1; }

    printf("Building ROSA index (%d conversations)...\n", fd_len(ds));
    rosa_index *idx = build_index(ds, tok);
    printf("Indexed %d tokens\n", idx->total_tokens);

    /* Load model */
    printf("Loading model from %s%s...\n", model_path, quantize ? " (INT8)" : "");
    rx_model *model = quantize ? rx_model_load_q8(model_path) : rx_model_load(model_path);
    if (!model) { fprintf(stderr, "ERROR: model\n"); return 1; }

    /* ---- Test questions ---- */
    typedef struct { const char *question; const char *hint; } qa_pair;
    qa_pair questions[] = {
        {"What is the WKV-7 kernel used for in RWKV-X?",
         "recurrent state evolution / generalized delta rule"},
        {"What does HybridMem use HDC binary vectors for?",
         "approximate similarity search / retrieval"},
        {"How many layers does the 0.2B RWKV-X model have?",
         "12 RWKV + 4 attention = 16 total"},
        {"What build flag enables BLAS acceleration in Foundry?",
         "USE_BLAS=1"},
        {"What is the purpose of LoRA adapters in the training system?",
         "fine-tune attention blocks / parameter-efficient"},
        {"What is the name chosen for the RWKV-X product?",
         "Engram"},
        {"What framework does Metal GPU use on macOS?",
         "Metal / MetalPerformanceShaders / Accelerate"},
        {"What format are model weights stored in?",
         "safetensors"},
        {"What is BootAI?",
         "bare-metal UEFI AI app / boots into LLM inference"},
        {"What tokenizer does RWKV-X use?",
         "RWKV World / BPE tokenizer"},
    };
    int n_questions = 10;

    printf("\n=== Phase 4: Model Answering (KEY TEST) ===\n\n");

    for (int q = 0; q < n_questions; q++) {
        printf("--- Question %d/%d ---\n", q + 1, n_questions);
        printf("Q: %s\n", questions[q].question);
        printf("Expected: %s\n\n", questions[q].hint);

        /* A) Baseline: no context */
        char baseline_prompt[2048];
        snprintf(baseline_prompt, sizeof(baseline_prompt),
            "<|user|>\n%s\n<|assistant|>\n", questions[q].question);

        printf("  [Baseline] generating...\n");
        clock_t t0 = clock();
        char *baseline_answer = generate_answer(model, tok, baseline_prompt, max_tokens);
        clock_t t1 = clock();
        printf("  Baseline (%.1fs): %.200s\n\n",
               (double)(t1 - t0) / CLOCKS_PER_SEC, baseline_answer);

        /* B) ROSA: retrieve context + question */
        char *context = retrieve_context(idx, tok, questions[q].question, context_window);

        char *rosa_prompt = (char *)malloc(context_window * 16 + 2048);
        if (context) {
            snprintf(rosa_prompt, context_window * 16 + 2048,
                "<|user|>\nContext from previous conversations:\n---\n%.4000s\n---\n\n"
                "Based on the context above, answer the following question:\n%s\n<|assistant|>\n",
                context, questions[q].question);
        } else {
            snprintf(rosa_prompt, 2048,
                "<|user|>\n%s\n<|assistant|>\n", questions[q].question);
        }

        printf("  [ROSA] generating (context: %s)...\n", context ? "yes" : "no");
        clock_t t2 = clock();
        char *rosa_answer = generate_answer(model, tok, rosa_prompt, max_tokens);
        clock_t t3 = clock();
        printf("  ROSA    (%.1fs): %.200s\n\n",
               (double)(t3 - t2) / CLOCKS_PER_SEC, rosa_answer);

        free(baseline_answer);
        free(rosa_answer);
        free(context);
        free(rosa_prompt);
    }

    printf("=== Phase 4 Complete ===\n");
    printf("Review answers above to score manually.\n");

    /* Cleanup */
    rx_model_free(model);
    free_index(idx);
    fd_free(ds);
    rx_tokenizer_free(tok);

    return 0;
}