/* retoor <retoor@molodetz.nl> */
#include "nfa.h"
#include <stdlib.h>
#include <string.h>

nfa_t *nfa_create(void) {
    nfa_t *nfa = malloc(sizeof(nfa_t));
    if (!nfa) return NULL;
    nfa->states = NULL;
    nfa->state_count = 0;
    nfa->capacity = 0;
    nfa->start = NULL;
    nfa->group_count = 0;
    return nfa;
}

void nfa_free(nfa_t *nfa) {
    if (!nfa) return;
    for (size_t i = 0; i < nfa->state_count; i++) {
        free(nfa->states[i]->transitions);
        free(nfa->states[i]);
    }
    free(nfa->states);
    free(nfa);
}

static bool nfa_grow(nfa_t *nfa) {
    size_t new_cap = nfa->capacity == 0 ? 16 : nfa->capacity * 2;
    if (new_cap > LOREG_MAX_STATES) {
        if (nfa->capacity >= LOREG_MAX_STATES) return false;
        new_cap = LOREG_MAX_STATES;
    }
    nfa_state_t **new_states = realloc(nfa->states, new_cap * sizeof(nfa_state_t *));
    if (!new_states) return false;
    nfa->states = new_states;
    nfa->capacity = new_cap;
    return true;
}

nfa_state_t *nfa_add_state(nfa_t *nfa) {
    if (nfa->state_count >= nfa->capacity) {
        if (!nfa_grow(nfa)) return NULL;
    }

    nfa_state_t *state = malloc(sizeof(nfa_state_t));
    if (!state) return NULL;

    state->id = (int)nfa->state_count;
    state->accepting = false;
    state->transitions = NULL;
    state->trans_count = 0;
    state->trans_capacity = 0;

    nfa->states[nfa->state_count++] = state;
    return state;
}

static bool transition_grow(nfa_state_t *state) {
    size_t new_cap = state->trans_capacity == 0 ? 4 : state->trans_capacity * 2;
    transition_t *new_trans = realloc(state->transitions, new_cap * sizeof(transition_t));
    if (!new_trans) return false;
    state->transitions = new_trans;
    state->trans_capacity = new_cap;
    return true;
}

void nfa_add_transition(nfa_state_t *from, nfa_state_t *to, transition_type_t type, char value) {
    if (from->trans_count >= from->trans_capacity) {
        if (!transition_grow(from)) return;
    }
    transition_t *t = &from->transitions[from->trans_count++];
    t->type = type;
    t->value = value;
    t->target = to;
    t->bracket = NULL;
    t->group_id = -1;
}

void nfa_add_bracket_transition(nfa_state_t *from, nfa_state_t *to, bracket_class_t *bracket) {
    if (from->trans_count >= from->trans_capacity) {
        if (!transition_grow(from)) return;
    }
    transition_t *t = &from->transitions[from->trans_count++];
    t->type = TRANS_BRACKET;
    t->value = '\0';
    t->target = to;
    t->bracket = bracket;
    t->group_id = -1;
}

void nfa_add_group_transition(nfa_state_t *from, nfa_state_t *to, transition_type_t type, int group_id) {
    if (from->trans_count >= from->trans_capacity) {
        if (!transition_grow(from)) return;
    }
    transition_t *t = &from->transitions[from->trans_count++];
    t->type = type;
    t->value = '\0';
    t->target = to;
    t->bracket = NULL;
    t->group_id = group_id;
}

static nfa_fragment_t build_nfa(nfa_t *nfa, ast_node_t *ast, loreg_error_t *error);

static nfa_fragment_t build_char(nfa_t *nfa, char c, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }
    nfa_add_transition(start, accept, TRANS_CHAR, c);
    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_dot(nfa_t *nfa, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }
    nfa_add_transition(start, accept, TRANS_DOT, '\0');
    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_class(nfa_t *nfa, transition_type_t type, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }
    nfa_add_transition(start, accept, type, '\0');
    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_bracket(nfa_t *nfa, bracket_class_t *bracket, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }
    nfa_add_bracket_transition(start, accept, bracket);
    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_concat(nfa_t *nfa, ast_node_t *left, ast_node_t *right, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_fragment_t left_frag = build_nfa(nfa, left, error);
    if (*error != LOREG_OK) return frag;
    nfa_fragment_t right_frag = build_nfa(nfa, right, error);
    if (*error != LOREG_OK) return frag;

    nfa_add_transition(left_frag.accept, right_frag.start, TRANS_EPSILON, '\0');
    frag.start = left_frag.start;
    frag.accept = right_frag.accept;
    return frag;
}

static nfa_fragment_t build_alter(nfa_t *nfa, ast_node_t *left, ast_node_t *right, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }

    nfa_fragment_t left_frag = build_nfa(nfa, left, error);
    if (*error != LOREG_OK) return frag;
    nfa_fragment_t right_frag = build_nfa(nfa, right, error);
    if (*error != LOREG_OK) return frag;

    nfa_add_transition(start, left_frag.start, TRANS_EPSILON, '\0');
    nfa_add_transition(start, right_frag.start, TRANS_EPSILON, '\0');
    nfa_add_transition(left_frag.accept, accept, TRANS_EPSILON, '\0');
    nfa_add_transition(right_frag.accept, accept, TRANS_EPSILON, '\0');

    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_star(nfa_t *nfa, ast_node_t *child, bool greedy, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }

    nfa_fragment_t child_frag = build_nfa(nfa, child, error);
    if (*error != LOREG_OK) return frag;

    if (greedy) {
        nfa_add_transition(start, child_frag.start, TRANS_EPSILON, '\0');
        nfa_add_transition(start, accept, TRANS_EPSILON, '\0');
    } else {
        nfa_add_transition(start, accept, TRANS_EPSILON, '\0');
        nfa_add_transition(start, child_frag.start, TRANS_EPSILON, '\0');
    }
    nfa_add_transition(child_frag.accept, child_frag.start, TRANS_EPSILON, '\0');
    nfa_add_transition(child_frag.accept, accept, TRANS_EPSILON, '\0');

    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_plus(nfa_t *nfa, ast_node_t *child, bool greedy, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }

    nfa_fragment_t child_frag = build_nfa(nfa, child, error);
    if (*error != LOREG_OK) return frag;

    if (greedy) {
        nfa_add_transition(child_frag.accept, child_frag.start, TRANS_EPSILON, '\0');
        nfa_add_transition(child_frag.accept, accept, TRANS_EPSILON, '\0');
    } else {
        nfa_add_transition(child_frag.accept, accept, TRANS_EPSILON, '\0');
        nfa_add_transition(child_frag.accept, child_frag.start, TRANS_EPSILON, '\0');
    }

    frag.start = child_frag.start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_question(nfa_t *nfa, ast_node_t *child, bool greedy, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }

    nfa_fragment_t child_frag = build_nfa(nfa, child, error);
    if (*error != LOREG_OK) return frag;

    if (greedy) {
        nfa_add_transition(start, child_frag.start, TRANS_EPSILON, '\0');
        nfa_add_transition(start, accept, TRANS_EPSILON, '\0');
    } else {
        nfa_add_transition(start, accept, TRANS_EPSILON, '\0');
        nfa_add_transition(start, child_frag.start, TRANS_EPSILON, '\0');
    }
    nfa_add_transition(child_frag.accept, accept, TRANS_EPSILON, '\0');

    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_group(nfa_t *nfa, ast_node_t *child, int group_id, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }

    nfa_fragment_t child_frag = build_nfa(nfa, child, error);
    if (*error != LOREG_OK) return frag;

    nfa_add_group_transition(start, child_frag.start, TRANS_GROUP_START, group_id);
    nfa_add_group_transition(child_frag.accept, accept, TRANS_GROUP_END, group_id);

    if (group_id + 1 > nfa->group_count) {
        nfa->group_count = group_id + 1;
    }

    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_anchor(nfa_t *nfa, transition_type_t type, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};
    nfa_state_t *start = nfa_add_state(nfa);
    nfa_state_t *accept = nfa_add_state(nfa);
    if (!start || !accept) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }
    nfa_add_transition(start, accept, type, '\0');
    frag.start = start;
    frag.accept = accept;
    return frag;
}

static nfa_fragment_t build_quantifier(nfa_t *nfa, ast_node_t *child, int min, int max, bool greedy, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};

    if (min == 0 && max == 0) {
        nfa_state_t *state = nfa_add_state(nfa);
        if (!state) {
            *error = LOREG_ERR_OUT_OF_MEMORY;
            return frag;
        }
        frag.start = state;
        frag.accept = state;
        return frag;
    }

    nfa_state_t *start = nfa_add_state(nfa);
    if (!start) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return frag;
    }

    nfa_state_t *current = start;

    for (int i = 0; i < min; i++) {
        nfa_fragment_t rep = build_nfa(nfa, child, error);
        if (*error != LOREG_OK) return frag;
        nfa_add_transition(current, rep.start, TRANS_EPSILON, '\0');
        current = rep.accept;
    }

    if (max < 0) {
        nfa_state_t *loop_start = nfa_add_state(nfa);
        nfa_state_t *accept = nfa_add_state(nfa);
        if (!loop_start || !accept) {
            *error = LOREG_ERR_OUT_OF_MEMORY;
            return frag;
        }

        nfa_add_transition(current, loop_start, TRANS_EPSILON, '\0');

        nfa_fragment_t rep = build_nfa(nfa, child, error);
        if (*error != LOREG_OK) return frag;

        if (greedy) {
            nfa_add_transition(loop_start, rep.start, TRANS_EPSILON, '\0');
            nfa_add_transition(loop_start, accept, TRANS_EPSILON, '\0');
        } else {
            nfa_add_transition(loop_start, accept, TRANS_EPSILON, '\0');
            nfa_add_transition(loop_start, rep.start, TRANS_EPSILON, '\0');
        }
        nfa_add_transition(rep.accept, loop_start, TRANS_EPSILON, '\0');

        frag.start = start;
        frag.accept = accept;
    } else {
        nfa_state_t *accept = nfa_add_state(nfa);
        if (!accept) {
            *error = LOREG_ERR_OUT_OF_MEMORY;
            return frag;
        }

        nfa_add_transition(current, accept, TRANS_EPSILON, '\0');

        for (int i = min; i < max; i++) {
            nfa_fragment_t rep = build_nfa(nfa, child, error);
            if (*error != LOREG_OK) return frag;

            if (greedy) {
                nfa_add_transition(current, rep.start, TRANS_EPSILON, '\0');
            } else {
                nfa_add_transition(current, accept, TRANS_EPSILON, '\0');
                nfa_add_transition(current, rep.start, TRANS_EPSILON, '\0');
            }

            if (greedy) {
                nfa_add_transition(rep.accept, accept, TRANS_EPSILON, '\0');
            }
            current = rep.accept;
        }

        if (!greedy) {
            nfa_add_transition(current, accept, TRANS_EPSILON, '\0');
        }

        frag.start = start;
        frag.accept = accept;
    }

    return frag;
}

static nfa_fragment_t build_nfa(nfa_t *nfa, ast_node_t *ast, loreg_error_t *error) {
    nfa_fragment_t frag = {NULL, NULL};

    if (!ast) {
        nfa_state_t *state = nfa_add_state(nfa);
        if (!state) {
            *error = LOREG_ERR_OUT_OF_MEMORY;
            return frag;
        }
        frag.start = state;
        frag.accept = state;
        return frag;
    }

    switch (ast->type) {
        case AST_CHAR:
            return build_char(nfa, ast->value, error);
        case AST_DOT:
            return build_dot(nfa, error);
        case AST_CONCAT:
            return build_concat(nfa, ast->left, ast->right, error);
        case AST_ALTER:
            return build_alter(nfa, ast->left, ast->right, error);
        case AST_STAR:
            return build_star(nfa, ast->left, ast->quant.greedy, error);
        case AST_PLUS:
            return build_plus(nfa, ast->left, ast->quant.greedy, error);
        case AST_QUESTION:
            return build_question(nfa, ast->left, ast->quant.greedy, error);
        case AST_GROUP:
            return build_group(nfa, ast->left, ast->group_id, error);
        case AST_ANCHOR_START:
            return build_anchor(nfa, TRANS_ANCHOR_START, error);
        case AST_ANCHOR_END:
            return build_anchor(nfa, TRANS_ANCHOR_END, error);
        case AST_BRACKET:
            return build_bracket(nfa, ast->bracket, error);
        case AST_QUANTIFIER:
            return build_quantifier(nfa, ast->left, ast->quant.min, ast->quant.max, ast->quant.greedy, error);
        case AST_CLASS_DIGIT:
            return build_class(nfa, TRANS_CLASS_DIGIT, error);
        case AST_CLASS_WORD:
            return build_class(nfa, TRANS_CLASS_WORD, error);
        case AST_CLASS_SPACE:
            return build_class(nfa, TRANS_CLASS_SPACE, error);
        case AST_CLASS_NDIGIT:
            return build_class(nfa, TRANS_CLASS_NDIGIT, error);
        case AST_CLASS_NWORD:
            return build_class(nfa, TRANS_CLASS_NWORD, error);
        case AST_CLASS_NSPACE:
            return build_class(nfa, TRANS_CLASS_NSPACE, error);
    }

    return frag;
}

nfa_t *nfa_from_ast(ast_node_t *ast, loreg_error_t *error) {
    *error = LOREG_OK;
    nfa_t *nfa = nfa_create();
    if (!nfa) {
        *error = LOREG_ERR_OUT_OF_MEMORY;
        return NULL;
    }

    nfa_fragment_t frag = build_nfa(nfa, ast, error);
    if (*error != LOREG_OK) {
        nfa_free(nfa);
        return NULL;
    }

    nfa->start = frag.start;
    frag.accept->accepting = true;

    return nfa;
}