/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * rosa.c — Suffix automaton implementation
 *
 * Binary automaton: fixed next[2] transitions.
 * Token automaton: open-addressing hash map transitions.
 */

#include "rosa/rosa.h"
#include <stdlib.h>
#include <string.h>

/* ================================================================
 * Binary suffix automaton
 * ================================================================ */

static void sa_init_state(sa_state *s) {
    s->len = 0;
    s->link = -1;
    s->endpos = -1;
    s->next[0] = -1;
    s->next[1] = -1;
}

static void sa_ensure_capacity(sa_automaton *sa, int needed) {
    if (needed <= sa->capacity) return;
    int new_cap = sa->capacity * 2;
    if (new_cap < needed) new_cap = needed;
    sa->states = (sa_state*)realloc(sa->states, new_cap * sizeof(sa_state));
    sa->capacity = new_cap;
}

sa_automaton* sa_new(int initial_cap) {
    if (initial_cap < 4) initial_cap = 4;
    sa_automaton *sa = (sa_automaton*)calloc(1, sizeof(sa_automaton));
    sa->states = (sa_state*)malloc(initial_cap * sizeof(sa_state));
    sa->capacity = initial_cap;
    sa->n_states = 1;
    sa->last = 0;
    sa->seq_len = 0;
    sa_init_state(&sa->states[0]);
    return sa;
}

void sa_free(sa_automaton *sa) {
    if (!sa) return;
    free(sa->states);
    free(sa);
}

void sa_reset(sa_automaton *sa) {
    if (!sa) return;
    sa->n_states = 1;
    sa->last = 0;
    sa->seq_len = 0;
    sa_init_state(&sa->states[0]);
}

void sa_extend(sa_automaton *sa, int symbol, int position) {
    sa_ensure_capacity(sa, sa->n_states + 2); /* extend may clone */

    int cur = sa->n_states++;
    sa_init_state(&sa->states[cur]);
    sa->states[cur].len = sa->states[sa->last].len + 1;
    sa->states[cur].endpos = position;

    int p = sa->last;
    while (p != -1 && sa->states[p].next[symbol] == -1) {
        sa->states[p].next[symbol] = cur;
        p = sa->states[p].link;
    }

    if (p == -1) {
        sa->states[cur].link = 0;
    } else {
        int q = sa->states[p].next[symbol];
        if (sa->states[p].len + 1 == sa->states[q].len) {
            sa->states[cur].link = q;
        } else {
            /* Clone q */
            sa_ensure_capacity(sa, sa->n_states + 1);
            int clone = sa->n_states++;
            sa->states[clone] = sa->states[q];
            sa->states[clone].len = sa->states[p].len + 1;

            while (p != -1 && sa->states[p].next[symbol] == q) {
                sa->states[p].next[symbol] = clone;
                p = sa->states[p].link;
            }
            sa->states[q].link = clone;
            sa->states[cur].link = clone;
        }
    }

    /* Update endpos along suffix link chain */
    int v = cur;
    while (v != -1 && sa->states[v].endpos < position) {
        sa->states[v].endpos = position;
        v = sa->states[v].link;
    }

    sa->last = cur;
    sa->seq_len++;
}

int sa_query_step(sa_automaton *sa, int symbol, int *state, int *match_len) {
    int v = *state;
    int len = *match_len;

    while (v != 0 && sa->states[v].next[symbol] == -1) {
        v = sa->states[v].link;
        len = sa->states[v].len;
    }

    if (sa->states[v].next[symbol] != -1) {
        v = sa->states[v].next[symbol];
        len++;
    }

    *state = v;
    *match_len = len;

    if (v != 0 && sa->states[v].endpos >= 0)
        return sa->states[v].endpos;
    return -1;
}

/* ================================================================
 * Token-level suffix automaton (hash-map transitions)
 * ================================================================ */

#define SA_TRANS_EMPTY (-1)
#define SA_TRANS_INIT_CAP 4

static void sa_trans_init(sa_token_state *s) {
    s->trans_cap = SA_TRANS_INIT_CAP;
    s->trans_count = 0;
    s->trans = (sa_trans_entry*)malloc(SA_TRANS_INIT_CAP * sizeof(sa_trans_entry));
    for (int i = 0; i < SA_TRANS_INIT_CAP; i++)
        s->trans[i].key = SA_TRANS_EMPTY;
}

static void sa_trans_free(sa_token_state *s) {
    free(s->trans);
    s->trans = NULL;
    s->trans_cap = 0;
    s->trans_count = 0;
}

static int sa_trans_find(const sa_token_state *s, int key) {
    int mask = s->trans_cap - 1;
    int idx = (unsigned)key & mask;
    for (int i = 0; i < s->trans_cap; i++) {
        int k = s->trans[(idx + i) & mask].key;
        if (k == key) return s->trans[(idx + i) & mask].value;
        if (k == SA_TRANS_EMPTY) return -1;
    }
    return -1;
}

static void sa_trans_rehash(sa_token_state *s) {
    int old_cap = s->trans_cap;
    sa_trans_entry *old = s->trans;
    s->trans_cap *= 2;
    s->trans = (sa_trans_entry*)malloc(s->trans_cap * sizeof(sa_trans_entry));
    for (int i = 0; i < s->trans_cap; i++)
        s->trans[i].key = SA_TRANS_EMPTY;
    int mask = s->trans_cap - 1;
    for (int i = 0; i < old_cap; i++) {
        if (old[i].key == SA_TRANS_EMPTY) continue;
        int idx = (unsigned)old[i].key & mask;
        while (s->trans[idx].key != SA_TRANS_EMPTY)
            idx = (idx + 1) & mask;
        s->trans[idx] = old[i];
    }
    free(old);
}

static void sa_trans_set(sa_token_state *s, int key, int value) {
    /* Check if key already exists */
    int mask = s->trans_cap - 1;
    int idx = (unsigned)key & mask;
    for (int i = 0; i < s->trans_cap; i++) {
        int pos = (idx + i) & mask;
        if (s->trans[pos].key == key) {
            s->trans[pos].value = value;
            return;
        }
        if (s->trans[pos].key == SA_TRANS_EMPTY) {
            s->trans[pos].key = key;
            s->trans[pos].value = value;
            s->trans_count++;
            if (s->trans_count * 2 > s->trans_cap)
                sa_trans_rehash(s);
            return;
        }
    }
    /* Table full (shouldn't happen with rehash) */
    sa_trans_rehash(s);
    sa_trans_set(s, key, value);
}

static void sa_trans_copy(sa_token_state *dst, const sa_token_state *src) {
    dst->trans_cap = src->trans_cap;
    dst->trans_count = src->trans_count;
    dst->trans = (sa_trans_entry*)malloc(dst->trans_cap * sizeof(sa_trans_entry));
    memcpy(dst->trans, src->trans, dst->trans_cap * sizeof(sa_trans_entry));
}

static void sa_token_init_state(sa_token_state *s) {
    s->len = 0;
    s->link = -1;
    s->endpos = -1;
    sa_trans_init(s);
}

static void sa_token_ensure_cap(sa_token_automaton *sa, int needed) {
    if (needed <= sa->capacity) return;
    int new_cap = sa->capacity * 2;
    if (new_cap < needed) new_cap = needed;
    sa->states = (sa_token_state*)realloc(sa->states, new_cap * sizeof(sa_token_state));
    sa->capacity = new_cap;
}

sa_token_automaton* sa_token_new(int initial_cap, int alphabet_size) {
    if (initial_cap < 4) initial_cap = 4;
    sa_token_automaton *sa = (sa_token_automaton*)calloc(1, sizeof(sa_token_automaton));
    sa->states = (sa_token_state*)malloc(initial_cap * sizeof(sa_token_state));
    sa->capacity = initial_cap;
    sa->n_states = 1;
    sa->last = 0;
    sa->seq_len = 0;
    sa->alphabet_size = alphabet_size;
    sa_token_init_state(&sa->states[0]);
    return sa;
}

void sa_token_free(sa_token_automaton *sa) {
    if (!sa) return;
    for (int i = 0; i < sa->n_states; i++)
        sa_trans_free(&sa->states[i]);
    free(sa->states);
    free(sa);
}

void sa_token_reset(sa_token_automaton *sa) {
    if (!sa) return;
    for (int i = 0; i < sa->n_states; i++)
        sa_trans_free(&sa->states[i]);
    sa->n_states = 1;
    sa->last = 0;
    sa->seq_len = 0;
    sa_token_init_state(&sa->states[0]);
}

void sa_token_extend(sa_token_automaton *sa, int token, int position) {
    sa_token_ensure_cap(sa, sa->n_states + 2);

    int cur = sa->n_states++;
    sa_token_init_state(&sa->states[cur]);
    sa->states[cur].len = sa->states[sa->last].len + 1;
    sa->states[cur].endpos = position;

    int p = sa->last;
    while (p != -1 && sa_trans_find(&sa->states[p], token) == -1) {
        sa_trans_set(&sa->states[p], token, cur);
        p = sa->states[p].link;
    }

    if (p == -1) {
        sa->states[cur].link = 0;
    } else {
        int q = sa_trans_find(&sa->states[p], token);
        if (sa->states[p].len + 1 == sa->states[q].len) {
            sa->states[cur].link = q;
        } else {
            sa_token_ensure_cap(sa, sa->n_states + 1);
            int clone = sa->n_states++;
            sa->states[clone].len = sa->states[p].len + 1;
            sa->states[clone].link = sa->states[q].link;
            sa->states[clone].endpos = sa->states[q].endpos;
            sa_trans_copy(&sa->states[clone], &sa->states[q]);

            while (p != -1 && sa_trans_find(&sa->states[p], token) == q) {
                sa_trans_set(&sa->states[p], token, clone);
                p = sa->states[p].link;
            }
            sa->states[q].link = clone;
            sa->states[cur].link = clone;
        }
    }

    int v = cur;
    while (v != -1 && sa->states[v].endpos < position) {
        sa->states[v].endpos = position;
        v = sa->states[v].link;
    }

    sa->last = cur;
    sa->seq_len++;
}

int sa_token_query_step(sa_token_automaton *sa, int token,
                        int *state, int *match_len) {
    int v = *state;
    int len = *match_len;

    while (v != 0 && sa_trans_find(&sa->states[v], token) == -1) {
        v = sa->states[v].link;
        len = sa->states[v].len;
    }

    int next = sa_trans_find(&sa->states[v], token);
    if (next != -1) {
        v = next;
        len++;
    }

    *state = v;
    *match_len = len;

    if (v != 0 && sa->states[v].endpos >= 0)
        return sa->states[v].endpos;
    return -1;
}

/* ================================================================
 * Position map
 * ================================================================ */

sa_position_map* sa_position_map_new(int initial_cap) {
    if (initial_cap < 16) initial_cap = 16;
    sa_position_map *map = (sa_position_map*)calloc(1, sizeof(sa_position_map));
    map->entries = (sa_position_info*)malloc(initial_cap * sizeof(sa_position_info));
    map->capacity = initial_cap;
    map->count = 0;
    return map;
}

void sa_position_map_free(sa_position_map *map) {
    if (!map) return;
    free(map->entries);
    free(map);
}

int sa_position_map_append(sa_position_map *map, int conv_idx,
                           int msg_idx, int token_offset) {
    if (map->count >= map->capacity) {
        map->capacity *= 2;
        map->entries = (sa_position_info*)realloc(map->entries,
            map->capacity * sizeof(sa_position_info));
    }
    int pos = map->count++;
    map->entries[pos].conv_idx = conv_idx;
    map->entries[pos].msg_idx = msg_idx;
    map->entries[pos].token_offset = token_offset;
    return pos;
}

const sa_position_info* sa_position_map_get(const sa_position_map *map,
                                             int position) {
    if (!map || position < 0 || position >= map->count) return NULL;
    return &map->entries[position];
}