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

	// Each TrieNode represents a single character in a word.
	// children[i] corresponds to character ('a' + i)
	// Assumes all input words contain only lowercase letters 'a' through 'z'
	typedef struct TrieNode {
		struct TrieNode * children[26];
		int isWord;
	} TrieNode;

	// Trie holds a continuous reference to the root
	// NOTE: Root is not representative of a letter
	// It simply enables branching towards other letters
	typedef struct Trie {
		TrieNode * root;
	} Trie;

	// Function Prototypes
	TrieNode * initTrieNode();
	Trie * initTrie();
	void addWord(Trie * trie, char * word);
	void addWordRecursive(TrieNode * node, char * word);
	int getWordCount(Trie * trie);
	int getWordCountRecursive(TrieNode * node);
	int getNodeCount(Trie * trie);
	int getNodeCountRecursive(TrieNode * node);
	int getWordCountWithPrefix(Trie * trie, char * prefix);
	void deleteWord(Trie * trie, char * word);
	int deleteWordRecursive(TrieNode * node, char * word);
	void emptyTrie(Trie * trie);
	void emptyTrieRecursive(TrieNode * node); 

	// Create and return initialized TrieNode
	TrieNode * initTrieNode() {
		TrieNode * node = malloc(sizeof(TrieNode));
		for (int i = 0; i < 26; i++) {
			node->children[i] = NULL;
		}
		node->isWord = 0;
		return node;
	}

	// Create and return initialized Trie struct with root created
	Trie * initTrie() {
		Trie * trie = malloc(sizeof(Trie));
		trie->root = initTrieNode();
		return trie;
	}

	// Wrapper function for addWordRecursive
	void addWord(Trie * trie, char * word) {
		if (!trie) return;
		addWordRecursive(trie->root, word);
	}

	// Recursively traverses through Trie by getting index of word[0] (word[0] - 'a')
	// node->children[word[0] - 'a'] points to node representing that character's node
	// If the node representing that letter is not in the Trie yet, that node is created and added
	// When recursively traversing towards next character, call word + 1 to shift starting index of 
	// word forward by one. i.e. "apple" + 1 => "pple"
	void addWordRecursive(TrieNode * node, char * word) {
		if (!node) return;
		
		if (word[0] == '\0') {
			node->isWord = 1;
			return;
		}
		
		int letterIdx = word[0] - 'a';  // maps 'a'->0, 'b'->1, ..., 'z'->25

		if (node->children[letterIdx] == NULL) 
			node->children[letterIdx] = initTrieNode();
		
		addWordRecursive(node->children[letterIdx], word + 1);
	}

	// Wrapper function for getWordCountRecursive
	int getWordCount(Trie * trie) {
		if (!trie) return 0;
		return getWordCountRecursive(trie->root);
	}

	// Counts all nodes in the trie in which isWord = 1
	int getWordCountRecursive(TrieNode * node) {
		if (!node) return 0;
		
		int count = node->isWord;
		
		for (int i = 0; i < 26; i++) 
			count += getWordCountRecursive(node->children[i]);

		return count; 
	}

	// Wrapper function for getNodeCountRecursive
	int getNodeCount(Trie * trie) {
		if (!trie) return 0;
		return getNodeCountRecursive(trie->root);
	}

	// Counts all nodes in the trie
	int getNodeCountRecursive(TrieNode * node) {
		if (!node) return 0;

		int count = 1;

		for (int i = 0; i < 26; i++) 
			count += getNodeCountRecursive(node->children[i]);

		return count; 
	}

	// Counts all words in the trie that start with a given prefix
	// i.e. "app", "apple", "appearance", and "appointment" all contain the prefix "app"
	int getWordCountWithPrefix(Trie * trie, char * prefix) {
		if (!trie) return 0;
		
		TrieNode * node = trie->root;
		while(prefix[0] != '\0') {
			int letterIdx = prefix[0] - 'a';
			if(!node->children[letterIdx]) return 0;
			node = node->children[letterIdx];
			prefix = prefix + 1;
		}

		return getWordCountRecursive(node);
	}

	// Wrapper function for deleteWordRecursive
	void deleteWord(Trie * trie, char * word) {
		if (!trie) return;
		deleteWordRecursive(trie->root, word);
	}

	// Deletes word from trie by setting isWord = 0 when end of word is reached
	// While recursively backtracking, nodes that do not have other dependencies are freed
	int deleteWordRecursive(TrieNode * node, char * word) {
		if (!node) return 0;

		if (word[0] == '\0') {
			if (!node->isWord) return 0;
			node->isWord = 0;

			for (int i = 0; i < 26; i++)
				if (node->children[i]) return 0;

			return 1; // delete this node
		}

		int letterIdx = word[0] - 'a';

		if (!node->children[letterIdx]) return 0;

		int shouldDeleteChild = deleteWordRecursive(node->children[letterIdx], word + 1);
	
		if (shouldDeleteChild) {
			free(node->children[letterIdx]);
			node->children[letterIdx] = NULL;

			if (node->isWord) return 0;

			for (int i = 0; i < 26; i++)
				if (node->children[i]) return 0;

			return 1;
		}

		return 0;
	}

	// Wrapper function for emptyTrieRecursive
	void emptyTrie(Trie * trie) {
		if(!trie) return;
		emptyTrieRecursive(trie->root);
		trie->root = NULL;
	}

	// Recursively traverses all nodes in trie and frees memory while backtracking
	void emptyTrieRecursive(TrieNode * node) { 
		if (!node) return;

		for (int i = 0; i < 26; i++) 
			emptyTrieRecursive(node->children[i]);

		free(node);
	}

	int main() {

		Trie * trie = initTrie();

		char *words[] = {"app","apple","strawberry","straw","stack",
			"appearance","appointment","strategy","tack","trend"};
		
		int numWords = 10;

		// Adds all words into the trie
		for (int i = 0; i < numWords; i++) {
			addWord(trie, words[i]);
		}

		// Prints word count: 10 and node count : 46
		printf("Current word count is: %d\n", getWordCount(trie)); 
		printf("Current node count is: %d\n\n", getNodeCount(trie));

		// Prints word count with prefix 'app': 4 and word count with prefix 'str': 3
		printf("Number of words beginning with 'app': %d\n", getWordCountWithPrefix(trie, "app"));
		printf("Number of words beginning with 'str': %d\n", getWordCountWithPrefix(trie, "str"));

		printf("\nDeleting 'appearance'...\n\n");
		deleteWord(trie, "appearance");

		// Prints word count: 9 and node count : 39
		printf("Current word count is: %d\n", getWordCount(trie));
		printf("Current node count is: %d\n\n", getNodeCount(trie));

		// Prints word count with prefix 'app': 3
		printf("Number of words beginning with 'app': %d\n\n", getWordCountWithPrefix(trie, "app"));

		printf("Destroying all nodes in trie...\n\n");
		emptyTrie(trie);

		// After emptyTrie, root is NULL, so both counts return 0
		printf("Current word count is: %d\n", getWordCount(trie));
		printf("Current node count is: %d\n\n", getNodeCount(trie));
		
		free(trie);

		return 0;
	}