/*
 * ahocorasick.c: implementation of ahocorasick library's functions
 * This file is part of multifast.
 *
 Copyright 2010-2012 Kamiar Kanani <kamiar.kanani@gmail.com>

 multifast is free software: you can redistribute it and/or modify
 it under the terms of the GNU Lesser General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.

 multifast is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU Lesser General Public License for more details.

 You should have received a copy of the GNU Lesser General Public License
 along with multifast.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

#include "ndpi_api.h"
#include "ahocorasick.h"

/* Allocation step for automata.all_nodes */
#define REALLOC_CHUNK_ALLNODES 200

/* Private function prototype */
static void ac_automata_register_nodeptr
(AC_AUTOMATA_t * thiz, AC_NODE_t * node);
static void ac_automata_union_matchstrs
(AC_NODE_t * node);
static void ac_automata_set_failure
(AC_AUTOMATA_t * thiz, AC_NODE_t * node, AC_ALPHABET_t * alphas);
static void ac_automata_traverse_setfailure
(AC_AUTOMATA_t * thiz, AC_NODE_t * node, AC_ALPHABET_t * alphas);


/******************************************************************************
 * FUNCTION: ac_automata_init
 * Initialize automata; allocate memories and set initial values
 * PARAMS:
 * MATCH_CALLBACK mc: call-back function
 * the call-back function will be used to reach the caller on match occurrence
 ******************************************************************************/
AC_AUTOMATA_t * ac_automata_init (MATCH_CALLBACK_f mc)
{
    AC_AUTOMATA_t * thiz = (AC_AUTOMATA_t *)ndpi_malloc(sizeof(AC_AUTOMATA_t));
    memset (thiz, 0, sizeof(AC_AUTOMATA_t));
    thiz->root = node_create ();
    thiz->all_nodes_max = REALLOC_CHUNK_ALLNODES;
    thiz->all_nodes = (AC_NODE_t **) ndpi_malloc (thiz->all_nodes_max*sizeof(AC_NODE_t *));
    thiz->match_callback = mc;
    ac_automata_register_nodeptr (thiz, thiz->root);
    ac_automata_reset (thiz);
    thiz->total_patterns = 0;
    thiz->automata_open = 1;
    return thiz;
}

/******************************************************************************
 * FUNCTION: ac_automata_add
 * Adds pattern to the automata.
 * PARAMS:
 * AC_AUTOMATA_t * thiz: the pointer to the automata
 * AC_PATTERN_t * patt: the pointer to added pattern
 * RETURN VALUE: AC_ERROR_t
 * the return value indicates the success or failure of adding action
 ******************************************************************************/
AC_ERROR_t ac_automata_add (AC_AUTOMATA_t * thiz, AC_PATTERN_t * patt)
{
    unsigned int i;
    AC_NODE_t * n = thiz->root;
    AC_NODE_t * next;
    AC_ALPHABET_t alpha;

    if(!thiz->automata_open)
        return ACERR_AUTOMATA_CLOSED;

    if (!patt->length)
        return ACERR_ZERO_PATTERN;

    if (patt->length > AC_PATTRN_MAX_LENGTH)
        return ACERR_LONG_PATTERN;

    for (i=0; i<patt->length; i++)
    {
        alpha = patt->astring[i];
        if ((next = node_find_next(n, alpha)))
        {
            n = next;
            continue;
        }
        else
        {
            next = node_create_next(n, alpha);
            next->depth = n->depth + 1;
            n = next;
            ac_automata_register_nodeptr(thiz, n);
        }
    }

    if(n->final)
        return ACERR_DUPLICATE_PATTERN;

    n->final = 1;
    node_register_matchstr(n, patt);
    thiz->total_patterns++;

    return ACERR_SUCCESS;
}

/******************************************************************************
 * FUNCTION: ac_automata_finalize
 * Locate the failure node for all nodes and collect all matched pattern for
 * every node. it also sorts outgoing edges of node, so binary search could be
 * performed on them. after calling this function the automate literally will
 * be finalized and you can not add new patterns to the automate.
 * PARAMS:
 * AC_AUTOMATA_t * thiz: the pointer to the automata
 ******************************************************************************/
void ac_automata_finalize (AC_AUTOMATA_t * thiz)
{
    unsigned int i;
    AC_ALPHABET_t *alphas;
    AC_NODE_t * node;

    if((alphas = ndpi_malloc(AC_PATTRN_MAX_LENGTH)) != NULL) {
        ac_automata_traverse_setfailure (thiz, thiz->root, alphas);

        for (i=0; i < thiz->all_nodes_num; i++)
        {
            node = thiz->all_nodes[i];
            ac_automata_union_matchstrs (node);
            node_sort_edges (node);
        }
        thiz->automata_open = 0; /* do not accept patterns any more */
        ndpi_free(alphas);
    }
}

/******************************************************************************
 * FUNCTION: ac_automata_search
 * Search in the input text using the given automata. on match event it will
 * call the call-back function. and the call-back function in turn after doing
 * its job, will return an integer value to ac_automata_search(). 0 value means
 * continue search, and non-0 value means stop search and return to the caller.
 * PARAMS:
 * AC_AUTOMATA_t * thiz: the pointer to the automata
 * AC_TEXT_t * txt: the input text that must be searched
 * void * param: this parameter will be send to call-back function. it is
 * useful for sending parameter to call-back function from caller function.
 * RETURN VALUE:
 * -1: failed call; automata is not finalized
 *  0: success; continue searching; call-back sent me a 0 value
 *  1: success; stop searching; call-back sent me a non-0 value
 ******************************************************************************/
int ac_automata_search (AC_AUTOMATA_t * thiz, AC_TEXT_t * txt, AC_REP_t * param)
{
    unsigned long position;
    AC_NODE_t *curr;
    AC_NODE_t *next;

    if(thiz->automata_open)
        /* you must call ac_automata_locate_failure() first */
        return -1;

    position = 0;
    curr = thiz->current_node;

    /* This is the main search loop.
     * it must be keep as lightweight as possible. */
    while (position < txt->length)
    {
        if(!(next = node_findbs_next(curr, txt->astring[position])))
        {
            if(curr->failure_node /* we are not in the root node */)
                curr = curr->failure_node;
            else
                position++;
        }
        else
        {
            curr = next;
            position++;
        }

        if(curr->final && next) {
            /* We check 'next' to find out if we came here after a alphabet
             * transition or due to a fail. in second case we should not report
             * matching because it was reported in previous node */
            thiz->match.position = position + thiz->base_position;
            thiz->match.match_num = curr->matched_patterns_num;
            thiz->match.patterns = curr->matched_patterns;
            /* we found a match! do call-back */
            if (thiz->match_callback(&thiz->match, txt, param))
                return 1;
        }
    }

    /* save status variables */
    thiz->current_node = curr;
    thiz->base_position += position;
    return 0;
}

/******************************************************************************
 * FUNCTION: ac_automata_reset
 * reset the automata and make it ready for doing new search on a new text.
 * when you finished with the input text, you must reset automata state for
 * new input, otherwise it will not work.
 * PARAMS:
 * AC_AUTOMATA_t * thiz: the pointer to the automata
 ******************************************************************************/
void ac_automata_reset (AC_AUTOMATA_t * thiz)
{
    thiz->current_node = thiz->root;
    thiz->base_position = 0;
}

/******************************************************************************
 * FUNCTION: ac_automata_release
 * Release all allocated memories to the automata
 * PARAMS:
 * AC_AUTOMATA_t * thiz: the pointer to the automata
 ******************************************************************************/
void ac_automata_release (AC_AUTOMATA_t * thiz)
{
    unsigned int i;
    AC_NODE_t * n;
    AC_PATTERN_t sid;

    for (i=0; i < thiz->all_nodes_num; i++)
    {
        n = thiz->all_nodes[i];

        if (n->matched_patterns_num) {
            sid = n->matched_patterns[0];

            if (sid.astring != NULL && sid.length > 0)
                ndpi_free(sid.astring);
        }

        node_release(n);
    }
    ndpi_free(thiz->all_nodes);
    ndpi_free(thiz);
}

/******************************************************************************
 * FUNCTION: ac_automata_display
 * Prints the automata to output in human readable form. it is useful for
 * debugging purpose.
 * PARAMS:
 * AC_AUTOMATA_t * thiz: the pointer to the automata
 * char repcast: 'n': print AC_REP_t as number, 's': print AC_REP_t as string
 ******************************************************************************/
void ac_automata_display (AC_AUTOMATA_t * thiz, char repcast)
{
    unsigned int i, j;
    AC_NODE_t * n;
    struct edge * e;
    AC_PATTERN_t sid;

    printf("---------------------------------\n");

    for (i=0; i<thiz->all_nodes_num; i++)
    {
        n = thiz->all_nodes[i];
        printf("NODE(%3d)/----fail----> NODE(%3d)\n",
                n->id, (n->failure_node)?n->failure_node->id:1);
        for (j=0; j<n->outgoing_degree; j++)
        {
            e = &n->outgoing[j];
            printf("         |----(");
            if(isgraph(e->alpha))
                printf("%c)---", e->alpha);
            else
                printf("0x%x)", e->alpha);
            printf("--> NODE(%3d)\n", e->next->id);
        }
        if (n->matched_patterns_num) {
            printf("Accepted patterns: {");
            for (j=0; j<n->matched_patterns_num; j++)
            {
                sid = n->matched_patterns[j];
                if(j) printf(", ");
                switch (repcast)
                {
                    case 'n':
                        printf("%u/%u/%u",
                                sid.rep.number,
                                sid.rep.category,
                                sid.rep.breed);
                        break;
                }
            }
            printf("}\n");
        }
        printf("---------------------------------\n");
    }
}

/******************************************************************************
 * FUNCTION: ac_automata_register_nodeptr
 * Adds the node pointer to all_nodes.
 ******************************************************************************/
static void ac_automata_register_nodeptr (AC_AUTOMATA_t * thiz, AC_NODE_t * node)
{
    if(thiz->all_nodes_num >= thiz->all_nodes_max)
    {
        thiz->all_nodes = ndpi_realloc(thiz->all_nodes, 
                thiz->all_nodes_max*sizeof(AC_NODE_t *),
                (REALLOC_CHUNK_ALLNODES+thiz->all_nodes_max)*sizeof(AC_NODE_t *)
                );
        thiz->all_nodes_max += REALLOC_CHUNK_ALLNODES;
    }
    thiz->all_nodes[thiz->all_nodes_num++] = node;
}

/******************************************************************************
 * FUNCTION: ac_automata_union_matchstrs
 * Collect accepted patterns of the node. the accepted patterns consist of the
 * node's own accepted pattern plus accepted patterns of its failure node.
 ******************************************************************************/
static void ac_automata_union_matchstrs (AC_NODE_t * node)
{
    unsigned int i;
    AC_NODE_t * m = node;

    while ((m = m->failure_node))
    {
        for (i=0; i < m->matched_patterns_num; i++)
            node_register_matchstr(node, &(m->matched_patterns[i]));

        if (m->final)
            node->final = 1;
    }
    // TODO : sort matched_patterns? is that necessary? I don't think so.
}

/******************************************************************************
 * FUNCTION: ac_automata_set_failure
 * find failure node for the given node.
 ******************************************************************************/
static void ac_automata_set_failure
(AC_AUTOMATA_t * thiz, AC_NODE_t * node, AC_ALPHABET_t * alphas)
{
    unsigned int i, j;
    AC_NODE_t * m;

    for (i=1; i < node->depth; i++)
    {
        m = thiz->root;
        for (j=i; j < node->depth && m; j++)
            m = node_find_next (m, alphas[j]);
        if (m)
        {
            node->failure_node = m;
            break;
        }
    }
    if (!node->failure_node)
        node->failure_node = thiz->root;
}

/******************************************************************************
 * FUNCTION: ac_automata_traverse_setfailure
 * Traverse all automata nodes using DFS (Depth First Search), meanwhile it set
 * the failure node for every node it passes through. this function must be
 * called after adding last pattern to automata. i.e. after calling this you
 * can not add further pattern to automata.
 ******************************************************************************/
static void ac_automata_traverse_setfailure
(AC_AUTOMATA_t * thiz, AC_NODE_t * node, AC_ALPHABET_t * alphas)
{
    unsigned int i;
    AC_NODE_t * next;

    for (i=0; i < node->outgoing_degree; i++)
    {
        alphas[node->depth] = node->outgoing[i].alpha;
        next = node->outgoing[i].next;

        /* At every node look for its failure node */
        ac_automata_set_failure (thiz, next, alphas);

        /* Recursively call itself to traverse all nodes */
        ac_automata_traverse_setfailure (thiz, next, alphas);
    }
}
