#pragma once

#include <cstdint>
#include <map>
#include <set>
#include <string>
#include <string_view>
#include <vector>

// Trie for matching multiple literals.
// This is used in common_peg_until_parser and to build a GBNF exclusion grammar
struct common_trie {
    struct node {
        std::map<uint32_t, size_t> children;  // Use uint32_t to store Unicode codepoints
        int32_t pattern = -1;                 // index of the pattern ending at this node, -1 if none
    };

    std::vector<node> nodes;

    common_trie() {
        create_node(); // root node
    }

    common_trie(const std::vector<std::string> & words) : common_trie() {
        for (const auto & w : words) {
            insert(w);
        }
    }

    enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH };

    // Check if a delimiter starts at the given position
    match_result check_at(std::string_view sv, size_t start_pos) const;

    // Insert a word as a sequence of Unicode codepoints, returns its pattern index
    int32_t insert(const std::string & word);

    // Insert a raw symbol sequence, returns its pattern index (insertion order,
    // duplicates keep the first index)
    int32_t insert(const std::vector<uint32_t> & symbols);

  private:
    int32_t n_patterns = 0;

    size_t create_node() {
        size_t index = nodes.size();
        nodes.emplace_back();
        return index;
    }
};

// Aho-Corasick automaton
struct common_aho_corasick {
    common_trie          t;
    std::vector<size_t>  fail;     // failure links
    std::vector<size_t>  order;    // states in BFS order
    std::vector<int32_t> match;    // longest pattern ending at each state (directly or via a suffix link), -1 if none
    std::set<uint32_t>   alphabet; // every character with a transition

    common_aho_corasick(common_trie trie);

    common_aho_corasick(const std::vector<std::string> & strings)
        : common_aho_corasick(common_trie(strings)) {}

    size_t num_states()          const { return t.nodes.size(); }
    bool   is_terminal(size_t s) const { return match[s] >= 0; }

    // index of the longest pattern ending at this state, -1 if none
    int32_t match_pattern(size_t s) const { return match[s]; }

    // follow failure links until a transition on `ch` exists.
    size_t next(size_t state, uint32_t ch) const;
};