package theorycrafter.utils


/**
 * A (node in the) prefix search trie.
 */
class SearchTrie<T> {


    /**
     * The items associated with this node.
     *
     * Performance note: this is optimized for adding; if there is a lot of removing going on, it's better to use a
     * multi-set to avoid a linear search when removing an association with an item.
     */
    private val searchItems = mutableListOf<T>()


    /**
     * The children of this node.
     */
    private val children = mutableMapOf<Char, SearchTrie<T>>()


    /**
     * Adds a search item, mapping it to the given text.
     *
     * Only queries of length at least [minQueryLength] will find this item. Note that the default value is 1, meaning
     * that by default, an empty query will find nothing.
     */
    fun add(item: T, text: String, minQueryLength: Int = 1) {
        var node = this
        var remainingQueryLength = minQueryLength

        for (index in text.indices){
            if (remainingQueryLength <= 0)
                node.searchItems.add(item)
            remainingQueryLength -= 1
            node = node.children.getOrPut(text[index], ::SearchTrie)
        }
        if (remainingQueryLength <= 0)
            node.searchItems.add(item)
    }


    /**
     * Removes the association of the given text with the given item.
     */
    fun remove(item: T, text: String) {
        remove(item, text, 0)
    }


    /**
     * Removes the association of the suffix of [text] starting at [index] with the given item.
     */
    private fun remove(item: T, text: String, index: Int) {
        searchItems.remove(item)  // It might actually not be present due to minQueryLength in add()

        if (index == text.length)
            return

        val child = children[text[index]]
            ?: throw IllegalArgumentException("Text \"$text\" is not in this SearchTrie")

        child.remove(item, text, index + 1)

        if (child.searchItems.isEmpty())
            children.remove(text[index])
    }



    /**
     * Returns all matches of the given word.
     */
    fun search(word: String): Iterable<T> {
        var tree: SearchTrie<T> = this
        for (c in word)
            tree = tree.children[c] ?: return emptySet()
        return tree.searchItems
    }


}
