/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_rosa_db_index.c — Phase 2: Index chat history with suffix automaton
 *
 * Loads conversations from SQLite, tokenizes them, builds a suffix automaton
 * over the full token stream, and creates a position map for retrieval.
 *
 * Usage:
 *   ./examples/test_rosa_db_index --vocab <vocab_path> --db <db_path>
 */

#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)  /* 64K tokens max per conversation */

static void print_usage(const char *prog) {
    fprintf(stderr, "Usage: %s --vocab <path> --db <path> [--max-conv N]\n", prog);
}

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

    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], "--max-conv") == 0 && i + 1 < argc)
            max_conv = atoi(argv[++i]);
        else {
            print_usage(argv[0]);
            return 1;
        }
    }

    if (!vocab_path || !db_path) {
        print_usage(argv[0]);
        return 1;
    }

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

    /* ---- Load conversations from SQLite ---- */
    printf("Loading conversations from %s...\n", db_path);
    fd_dataset *ds = fd_load_chat_sqlite(db_path, max_conv);
    if (!ds) {
        fprintf(stderr, "ERROR: Failed to load database\n");
        rx_tokenizer_free(tok);
        return 1;
    }

    int n_conv = fd_len(ds);
    printf("Loaded %d conversations\n\n", n_conv);

    /* ---- Tokenize all conversations and build suffix automaton ---- */
    printf("Tokenizing and indexing...\n");

    uint32_t *token_buf = (uint32_t *)malloc(MAX_TOKENS_PER_CONV * sizeof(uint32_t));
    sa_token_automaton *sa = sa_token_new(1024, 65536);
    sa_position_map *pmap = sa_position_map_new(4096);

    int total_tokens = 0;
    int total_messages = 0;
    clock_t t_start = clock();

    for (int c = 0; c < n_conv; c++) {
        fd_sample *sample = fd_get(ds, c);
        const char *text = fd_sample_get(sample, "text");
        if (!text) continue;

        /* Count messages (approximate: count <|user|> and <|assistant|> tags) */
        int msg_count = 0;
        const char *p = text;
        while ((p = strstr(p, "<|")) != NULL) {
            if (strncmp(p, "<|user|>", 8) == 0 || strncmp(p, "<|assistant|>", 13) == 0)
                msg_count++;
            p += 2;
        }
        total_messages += msg_count;

        /* Tokenize the full conversation text */
        int n_tokens = rx_tokenizer_encode(tok, text, token_buf, MAX_TOKENS_PER_CONV);
        if (n_tokens <= 0) continue;

        /* Extend suffix automaton and position map for each token */
        for (int t = 0; t < n_tokens; t++) {
            int global_pos = total_tokens + t;
            sa_token_extend(sa, (int)token_buf[t], global_pos);
            sa_position_map_append(pmap, c, 0, t);
        }

        total_tokens += n_tokens;

        /* Progress every 5 conversations */
        if ((c + 1) % 5 == 0 || c == n_conv - 1) {
            printf("  [%d/%d] %d tokens, %d SA states\r",
                   c + 1, n_conv, total_tokens, sa->n_states);
            fflush(stdout);
        }
    }

    clock_t t_end = clock();
    double index_time = (double)(t_end - t_start) / CLOCKS_PER_SEC;

    printf("\n\n");

    /* ---- Compute memory usage ---- */
    /* SA memory: n_states * sizeof(sa_token_state) + transitions */
    size_t sa_mem = (size_t)sa->n_states * sizeof(sa_token_state);
    for (int i = 0; i < sa->n_states; i++)
        sa_mem += (size_t)sa->states[i].trans_cap * sizeof(sa_trans_entry);
    double sa_mem_mb = (double)sa_mem / (1024.0 * 1024.0);

    /* Position map memory */
    size_t pmap_mem = (size_t)pmap->count * sizeof(sa_position_info);
    double pmap_mem_mb = (double)pmap_mem / (1024.0 * 1024.0);

    /* ---- Verify SA property: states <= 2*N + 1 ---- */
    int sa_limit = 2 * total_tokens + 1;
    int sa_ok = sa->n_states <= sa_limit;

    /* ---- Quick query test ---- */
    printf("Running quick query test...\n");
    int query_state = 0, query_match = 0;
    int query_endpos = -1;

    /* Query with the first 10 tokens of the first conversation */
    fd_sample *first = fd_get(ds, 0);
    const char *first_text = fd_sample_get(first, "text");
    int first_n = rx_tokenizer_encode(tok, first_text, token_buf, MAX_TOKENS_PER_CONV);
    int query_len = first_n < 10 ? first_n : 10;

    for (int t = 0; t < query_len; t++)
        query_endpos = sa_token_query_step(sa, (int)token_buf[t], &query_state, &query_match);

    printf("  Query: first %d tokens of conv 0\n", query_len);
    printf("  Match length: %d tokens\n", query_match);
    printf("  End position: %d\n", query_endpos);
    if (query_endpos >= 0 && query_endpos < pmap->count) {
        const sa_position_info *info = sa_position_map_get(pmap, query_endpos);
        printf("  Mapped to: conv %d, token offset %d\n", info->conv_idx, info->token_offset);
    }
    printf("\n");

    /* ---- Print results ---- */
    printf("=== Phase 2 Results ===\n\n");
    printf("| Metric                  | Value          |\n");
    printf("|-------------------------|----------------|\n");
    printf("| Conversations loaded    | %d             |\n", n_conv);
    printf("| Messages loaded         | %d             |\n", total_messages);
    printf("| Total tokens            | %d             |\n", total_tokens);
    printf("| SA states               | %d             |\n", sa->n_states);
    printf("| SA states <= 2*N+1      | %s             |\n", sa_ok ? "YES" : "NO");
    printf("| SA memory (MB)          | %.2f           |\n", sa_mem_mb);
    printf("| Position map memory (MB)| %.2f           |\n", pmap_mem_mb);
    printf("| Total memory (MB)       | %.2f           |\n", sa_mem_mb + pmap_mem_mb);
    printf("| Index time (s)          | %.2f           |\n", index_time);
    printf("\n");

    /* ---- Pass/fail ---- */
    int pass = 1;
    if (!sa_ok) {
        printf("FAIL: SA states (%d) > 2*N+1 (%d)\n", sa->n_states, sa_limit);
        pass = 0;
    }
    if (sa_mem_mb + pmap_mem_mb > 500.0) {
        printf("FAIL: Memory %.2f MB > 500 MB limit\n", sa_mem_mb + pmap_mem_mb);
        pass = 0;
    }
    if (index_time > 30.0) {
        printf("FAIL: Index time %.2f s > 30 s limit\n", index_time);
        pass = 0;
    }
    if (query_match < query_len) {
        printf("FAIL: Query self-match %d < %d (should match own tokens)\n", query_match, query_len);
        pass = 0;
    }

    printf("**Verdict: %s**\n", pass ? "PASS" : "FAIL");

    /* ---- Cleanup ---- */
    sa_token_free(sa);
    sa_position_map_free(pmap);
    free(token_buf);
    fd_free(ds);
    rx_tokenizer_free(tok);

    return pass ? 0 : 1;
}