/*
 * PrefixTree.kt
 * 
 * This is Dave Musicant's Trie.kt from Fall 2024,
 * which was ported by Dave from the Trie code at
 * https://www.geeksforgeeks.org/introduction-to-trie-data-structure-and-algorithm-tutorials.
 * 
 * Modified lightly by Tanya Amert for Spring 2025.
 * 
 * Illustrates an implementation of a Prefix Tree ("Trie").
 */

 data class TrieNode(
    var childNode: Array<TrieNode?> = Array<TrieNode?>(26) { null },

    // This will keep track of number of strings that are stored in
    // the prefix tree from root node to any Trie node.
    var wordCount: Int = 0
)

/*
 * Add a string to the prefix tree, character by character if needed.
 */
fun insert(root: TrieNode, key: String) {
    // Initialize the currentNode pointer with the root node
    var currentNode = root

    // Note that letter is a 'Char' type
    for (letter in key) {
        val index = letter - 'a'

        // Check if the node exist for the current character in the Trie
        if (currentNode.childNode[index] == null) {

            // Keep the reference for the newly created node
            currentNode.childNode[index] = TrieNode()
        }

        // Now, move the current node pointer to the newly created node
        currentNode = currentNode.childNode[index]!!
    }

    // Increment the wordEndCount for the last currentNode pointer;
    // this implies that there is a string ending at currentNode
    currentNode.wordCount++
}

/*
 * Checks if a given prefix exists in the prefix tree.
 */
fun isPrefixExist(root: TrieNode, key: String): Boolean {
    // Initialize the currentNode pointer with the root node
    var currentNode = root

    // Iterate across the length of the string
    for (c in key) {

        // Check if the node exists for the current character in the Trie
        if (currentNode.childNode[c - 'a'] == null) {

            // Given word as a prefix does not exist in Trie
            return false
        }

        // Move the currentNode pointer to the already
        // existing node for current character
        currentNode = currentNode.childNode[c - 'a']!!
    }

    // Prefix exists in the Trie
    return true
}

/*
 * Checks whether a given string is present in the prefix tree.
 */
fun search(root: TrieNode, key: String): Boolean {
    // Initialize the currentNode with the root node
    var currentNode = root

    for (letter in key) {
        val index = letter - 'a'

        // Check if the node exist for the current character in the Trie
        if (currentNode.childNode[index] == null) {
            return false
        }

        // Move the currentNode to the already existing node for current
        // character
        currentNode = currentNode.childNode[index]!!
    }

    return (currentNode.wordCount > 0)
}

/*
 * Deletes a string from the prefix tree, clearing out pointers to
 * lower nodes if a node is no longer used.
 */
fun deleteKey(root: TrieNode, word: String): Boolean {
    var currentNode = root
    var lastBranchNode: TrieNode? = null
    var lastBranchChar = 'a'

    for (c in word) {
        if (currentNode.childNode[c - 'a'] == null) {
            // If the current node has no child, the word is not present
            return false
        } else {
            var count = 0
            // Count the number of non-null child nodes
            for (i in 0..<26) {
                if (currentNode.childNode[i] != null)
                    count++
            }

            if (count > 1) {
                // If there is more than one child, store the last branch information
                lastBranchNode = currentNode
                lastBranchChar = c
            }
            currentNode = currentNode.childNode[c - 'a']!!
        }
    }

    var count = 0
    // Count the number of non-null child nodes at the last character
    for (i in 0..<26) {
        if (currentNode.childNode[i] != null)
            count++
    }

    // Case 1: the deleted word is a prefix of other words in Trie
    if (count > 0) {
        // Decrement the word count and indicate successful deletion
        currentNode.wordCount--
        return true
    }

    // Case 2: the deleted word shares a common prefix with other
    // words in Trie
    if (lastBranchNode != null) {
        // Remove the link to the deleted word
        lastBranchNode.childNode[lastBranchChar - 'a'] = null
        return true
    }
    // Case 3: The deleted word does not share any common prefix with
    // other words in Trie
    else {
        // Remove the link to the deleted word from the root
        root.childNode[word[0] - 'a'] = null
        return true
    }
}

fun main() {
    // Make a root node for the prefix tree
    val root = TrieNode()

    // List out the strings that we want to *insert* in the prefix tree
    val inputStrings = listOf("and", "ant", "do", "geek", "dad", "ball", "an")

    for (inputString in inputStrings) {
        insert(root, inputString)
    }

    // List out the strings that we want to *search* in the prefix tree
    val searchQueryStrings = listOf("do", "geek", "bat", "an")

    for (searchQueryString in searchQueryStrings) {
        println("Query String: $searchQueryString")
        if (search(root, searchQueryString)) {
            println("The query string is present in the Trie")
        }
        else {
            println("The query string is not present in the Trie")
        }
    }

    // List out the strings that we want to *delete* from the prefix tree
    val deleteQueryStrings = listOf("geek", "tea")

    for (deleteQueryString in deleteQueryStrings) {
        println("Query String: $deleteQueryString")
        if (deleteKey(root, deleteQueryString)) {
            println("The query string is successfully deleted")
        }
        else {
            println("The query string is not present in the Trie")
        }
    }
}