#include "WordTree.hpp"

WordTree::WordTree() :
    m_root(std::make_shared<TreeNode>()) {}

WordTree::~WordTree()
{
}

WordTree::TreeNode::TreeNode() :
    m_endOfWord(false), m_children()
{
}

void WordTree::add(const std::string& word)
{
    if (word.empty())
        return;

    std::shared_ptr<TreeNode> current = m_root;

    for (char c : word)
    {
        if (current->m_children[c - 'a'] == nullptr)
            current->m_children[c - 'a'] = std::make_shared<TreeNode>();
        current = current->m_children[c - 'a'];
    }
    current->m_endOfWord = true;
}

bool WordTree::find(const std::string& word)
{
    std::shared_ptr<TreeNode> current = m_root;

    for (char c : word)
    {
        if (current->m_children[c - 'a'] == nullptr)
            return false;
        current = current->m_children[c - 'a'];
    }
    return current->m_endOfWord;
}

std::vector<std::string> WordTree::predict(std::string partial, std::uint8_t howMany)
{
    std::vector<std::string> predictions;
    if (partial.empty())
        return predictions;

    std::shared_ptr<TreeNode> partialEnd = m_root;

    for (char c : partial)
    {
        if (partialEnd->m_children[c - 'a'] == nullptr)
            return predictions;
        partialEnd = partialEnd->m_children[c - 'a'];
    }

    std::queue<std::pair<std::string, std::shared_ptr<TreeNode>>> queue;
    queue.push({ partial, partialEnd });

    while (!queue.empty() && predictions.size() < howMany)
    {
        auto curr = queue.front();
        queue.pop();

        if (curr.second->m_endOfWord && curr.second != partialEnd)
            predictions.push_back(curr.first);

        for (int i = 0; i < 26; i++)
            if (curr.second->m_children[i] != nullptr)
                queue.push({ curr.first + static_cast<char>(i + 'a'), curr.second->m_children[i] });
    }

    return predictions;
}

std::size_t WordTree::size()
{
    std::size_t size = 0;

    std::queue<std::shared_ptr<TreeNode>> queue;
    queue.push(m_root);

    while (!queue.empty())
    {
        auto curr = queue.front();
        queue.pop();

        if (curr->m_endOfWord)
            size++;

        for (int i = 0; i < 26; i++)
            if (curr->m_children[i] != nullptr)
                queue.push(curr->m_children[i]);
    }
    return size;
}