#include "Parsodus/util/firstset.h"
#include "Parsodus/grammar.h"

namespace pds {
namespace util {

FirstSet::FirstSet(const Grammar& g)
{
    for (std::string term : g.terminals) {
        m_first[term].insert(term);
    }

    bool changes = true;

    auto update = [&changes, this](std::string head, const std::set<std::string>& s) {
        for (std::string elem : s) {
            if (!m_first[head].count(elem)) {
                changes = true;
                m_first[head].insert(s.begin(), s.end());
                return;
            }
        }
    };

    while (changes) {
        changes = false;

        for (const auto& rule : g.rules) {
            if (rule->tail.size() == 0) {
                update(rule->head, {""});
            }

            bool ended = false;
            const std::string& head = rule->head;
            for (const std::string& replacement : rule->tail) {
                if (m_first[replacement].count("")) {
                    std::set<std::string> tmp = m_first[replacement];
                    tmp.erase("");
                    update(head, tmp);
                } else {
                    update(head, m_first[replacement]);
                    ended = true;
                    break;
                }
            }
            if (!ended) {
                update(head, {""});
            }
        }
    }
}

std::set<std::string> FirstSet::operator()(std::string key) const {
    auto tmp = m_first.find(key);
    if (tmp == m_first.end())
        return {};
    return tmp->second;
}

std::set<std::string> FirstSet::operator()(std::vector<std::string> sequence) const {
    std::set<std::string> result;
    bool hasEmpty = false;
    for (std::string key : sequence) {
        std::set<std::string> tmp = (*this)(key);
        hasEmpty = tmp.count("");
        tmp.erase("");
        result.insert(tmp.begin(), tmp.end());
        if (!hasEmpty)
            break;
    }
    if (hasEmpty)
        result.insert("");

    return result;
}

/*
 * template <typename Item>
 * void Generator<Item>::buildFollow() {
 *     //EOF follow the added start rule.
 *     m_follow[EXTENDED_START].insert(EOF_PLACEHOLDER);
 * 
 *     bool changes = true;
 * 
 *     auto update = [&changes, this](std::string head, auto s) {
 *         s.erase("");
 *         for (std::string elem : s) {
 *             if (!m_follow[head].count(elem)) {
 *                 changes = true;
 *                 m_follow[head].insert(s.begin(), s.end());
 *                 return;
 *             }
 *         }
 *     };
 * 
 *     while (changes) {
 *         changes = false;
 *         for (const auto& p : m_gram.rules) {
 *             std::string& head = p.first;
 *             for (const auto& rule : p.second) {
 *                 for (std::size_t i = 0; i < rule.tail.size(); i++) {
 *                     if (!m_gram.terminals.count(rule.tail[i])) {
 *                         if (i == rule.tail.size() - 1 || m_first[rule.tail[i + 1]].count("")) {
 *                             update(rule.tail[i], m_follow[head]);
 *                         }
 *                         if (i < rule.tail.size() - 1) {
 *                             update(rule.tail[i], m_first[rule.tail[i + 1]]);
 *                         }
 *                     }
 *                 }
 *             }
 *         }
 *     }
 * }
 */

} /* util */
} /* pds */