/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_rosa_db_retrieve.c — Phase 3: Retrieval quality test
 *
 * Indexes all conversations, then runs 10 test queries.
 * For each query, finds the longest suffix match and retrieves
 * a context window around the match point.
 *
 * Usage:
 *   ./examples/test_rosa_db_retrieve \
 *     --vocab <vocab_path> --db <db_path> --context-window 256
 */

#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

/* ---- Index builder (same as Phase 2) ---- */

typedef struct {
    sa_token_automaton *sa;
    sa_position_map *pmap;
    int total_tokens;
    int n_conv;
    /* Per-conversation token offsets for decoding context */
    int *conv_start;   /* global position where each conv starts */
    int *conv_len;     /* token count per conv */
    uint32_t *all_tokens; /* full token stream for decoding */
    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;

        /* Grow all_tokens if needed */
        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);
}

/* ---- Query and retrieve context ---- */

typedef struct {
    int match_len;
    int endpos;
    int conv_idx;
    int token_offset;
    char *context;       /* decoded context text (malloc'd) */
    double query_time_ms;
} retrieve_result;

static retrieve_result query_and_retrieve(rosa_index *idx, rx_tokenizer *tok,
                                           const char *query_text, int context_window) {
    retrieve_result res = {0};
    res.endpos = -1;
    res.conv_idx = -1;

    /* Tokenize query */
    uint32_t query_tokens[MAX_QUERY_TOKENS];
    int qlen = rx_tokenizer_encode(tok, query_text, query_tokens, MAX_QUERY_TOKENS);
    if (qlen <= 0) return res;

    /* Run query through SA */
    clock_t t0 = clock();
    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)query_tokens[i], &state, &match_len);
        if (match_len > best_match) {
            best_match = match_len;
            best_endpos = ep;
        }
    }
    clock_t t1 = clock();
    res.query_time_ms = (double)(t1 - t0) / CLOCKS_PER_SEC * 1000.0;
    res.match_len = best_match;
    res.endpos = best_endpos;

    if (best_endpos < 0) return res;

    /* Look up position */
    const sa_position_info *info = sa_position_map_get(idx->pmap, best_endpos);
    if (!info) return res;
    res.conv_idx = info->conv_idx;
    res.token_offset = info->token_offset;

    /* Extract context window around the match */
    int conv_start = idx->conv_start[res.conv_idx];
    int conv_len = idx->conv_len[res.conv_idx];

    /* Match ends at token_offset within the conversation */
    int match_end_in_conv = info->token_offset;
    int ctx_start = match_end_in_conv - best_match - context_window / 2;
    if (ctx_start < 0) ctx_start = 0;
    int ctx_end = match_end_in_conv + context_window / 2;
    if (ctx_end > conv_len) ctx_end = conv_len;

    int ctx_len = ctx_end - ctx_start;
    if (ctx_len > 0) {
        res.context = rx_tokenizer_decode(tok,
            &idx->all_tokens[conv_start + ctx_start], ctx_len);
    }

    return res;
}

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

int main(int argc, char **argv) {
    const char *vocab_path = NULL;
    const char *db_path = NULL;
    int context_window = 256;

    for (int i = 1; i < argc; i++) {
        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]);
    }

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

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

    printf("Loading conversations...\n");
    fd_dataset *ds = fd_load_chat_sqlite(db_path, 0);
    if (!ds) { fprintf(stderr, "ERROR: database\n"); return 1; }
    printf("Loaded %d conversations\n", fd_len(ds));

    /* Build index */
    printf("Building ROSA index...\n");
    clock_t t0 = clock();
    rosa_index *idx = build_index(ds, tok);
    clock_t t1 = clock();
    printf("Indexed %d tokens in %.2f s\n\n", idx->total_tokens,
           (double)(t1 - t0) / CLOCKS_PER_SEC);

    /* ---- 10 test queries ---- */
    const char *queries[] = {
        "suffix automaton retrieval",
        "Metal GPU acceleration kernel",
        "RWKV-7 WKV kernel implementation",
        "HybridMem HDC binary vectors",
        "LoRA training attention blocks",
        "Apple Neural Engine ANE CoreML",
        "INT8 quantization NEON dot product",
        "DeepEmbed per-token embedding",
        "UEFI boot application bare metal",
        "tokenizer BPE encode decode vocabulary",
    };
    int n_queries = 10;

    printf("=== Phase 3: Retrieval Quality ===\n\n");
    printf("Context window: %d tokens\n\n", context_window);

    printf("| # | Query | Match (tok) | Conv | Time (ms) | Context snippet |\n");
    printf("|---|-------|-------------|------|-----------|------------------|\n");

    int total_matched = 0;

    for (int q = 0; q < n_queries; q++) {
        retrieve_result res = query_and_retrieve(idx, tok, queries[q], context_window);

        /* Truncate context for display */
        char snippet[201] = "(no match)";
        if (res.context) {
            /* Find first interesting part (skip whitespace/tags) */
            int slen = (int)strlen(res.context);
            int start = 0;
            if (slen > 200) {
                /* Try to start near the match point */
                start = slen / 2 - 100;
                if (start < 0) start = 0;
            }
            int copy_len = slen - start;
            if (copy_len > 200) copy_len = 200;
            memcpy(snippet, res.context + start, copy_len);
            snippet[copy_len] = '\0';
            /* Replace newlines with spaces for table display */
            for (int i = 0; i < copy_len; i++)
                if (snippet[i] == '\n' || snippet[i] == '\r') snippet[i] = ' ';
            /* Replace pipes for markdown table */
            for (int i = 0; i < copy_len; i++)
                if (snippet[i] == '|') snippet[i] = '/';
        }

        if (res.match_len > 0) total_matched++;

        printf("| %d | %s | %d | %d | %.2f | %.80s... |\n",
               q + 1, queries[q], res.match_len,
               res.conv_idx, res.query_time_ms, snippet);

        free(res.context);
    }

    printf("\n**Queries with matches: %d/%d**\n", total_matched, n_queries);
    printf("\n**Verdict: %s**\n", total_matched >= 7 ? "PASS" : (total_matched >= 4 ? "MARGINAL" : "FAIL"));

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

    return total_matched >= 4 ? 0 : 1;
}