// Arup Guha
// 4/25/2024
// Solution to Spring 2024 COP 3502H Final Exam Question 13

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

typedef struct trie {
    int isWord;
    struct trie* next[26];
} trie;

void insert(trie* root, char* str, int k, int n);
trie* mkNode();
int* wordlengths(trie* root, char* str, int* ptrNumWords);
void freetrie(trie* root);

int main(void) {

    // Make a small trie.
    trie* mine = mkNode();
    insert(mine, "sells", 0, 5);
    insert(mine, "sea", 0, 3);
    insert(mine, "shells", 0, 6);
    insert(mine, "shore", 0, 5);
    insert(mine, "by", 0, 2);
    insert(mine, "sally", 0, 5);
    insert(mine, "the", 0, 3);

    // Run it.
    int numWords;
    int* len = wordlengths(mine, "sallysellsseashellsbytheseashore", &numWords);

    // Print lengths, should be 5,5,3,6,2,3,3,5
    for (int i=0; i<numWords; i++)
        printf("%d ", len[i]);

    // Free stuff.
    free(len);
    freetrie(mine);

    return 0;
}

void freetrie(trie* root) {
    if (root == NULL) return;
    for (int i=0; i<26; i++)
        freetrie(root->next[i]);
    free(root);
}

// Returns an empty trie node.
trie* mkNode() {
    trie* res = malloc(sizeof(trie));
    res->isWord = 0;
    for (int i=0; i<26; i++)
        res->next[i]= NULL;
    return res;
}
void insert(trie* root, char* str, int k, int n) {

    // Made it to end.
    if (k == n) {
        root->isWord = 1;
        return;
    }

    // Make node if it's not there.
    if (root->next[str[k]-'a'] == NULL)
        root->next[str[k]-'a'] = mkNode();

    // recurse.
    insert(root->next[str[k]-'a'], str, k+1, n);
}

int* wordlengths(trie* root, char* str, int* ptrNumWords) {

    // To make easier to type and more efficient.
    int n = strlen(str);
    int* res = calloc(n, sizeof(int));

    // i is index into str, j is word number we are on, 0-based.
    int i = 0, j = 0;

    // Loop through full string.
    while (i<n) {

        // Beginning of a word.
        trie* tmp = root;
        int oldi = i;

        // Due to prefix rule, we can just wait till we hit a word.
        while (!tmp->isWord) {

            // Follow the appropriate link and go to the next letter.
            tmp = tmp->next[str[i]-'a'];
            i++;
        }

        // This is the length of the word we just finished reading in.
        res[j] = i - oldi;

        // Advance to next word.
        j++;
    }

    // Update the size of this array to be correct.
    res = realloc(res, j*sizeof(int));

    // Store the size of this array.
    *ptrNumWords = j;

    // Return a pointer to the array.
    return res;
}