#include "Lexesis/backends/cpp.h"

#include <cassert>
#include <iostream>

namespace {
    //Some shortcut utility functions for creating a TemplateContext
    
    lxs::templ::TemplateContext make_map_elem(std::string key, std::string value) {
        return lxs::templ::make_map({{key, lxs::templ::make_string(value)}});
    }

    lxs::templ::TemplateContext make_int(int i) {
        return lxs::templ::make_string(std::to_string(i));
    }

    lxs::templ::TemplateContext make_State(lxs::State i) {
        return lxs::templ::make_string(std::to_string(i));
    }
}

namespace lxs { namespace backends {

    CppBackend::CppBackend() : Backend()
    {}

    CppBackend::~CppBackend()
    {}

    std::string CppBackend::getName() {
        return "c++";
    }

    bool CppBackend::canProcessLang(std::string lang) {
        for (char& c : lang)
            c = std::tolower(c);
        return lang == "c++" || lang == "cpp" || lang == "cxx";
    }

    void CppBackend::generateLexer(
            std::function<std::unique_ptr<std::ostream>(std::string)> getOstreamForFileName,
            std::string lexerName,
            const DFA& dfa)
    {
        assert(lexerName.length());

        lexerName[0] = std::toupper(lexerName[0]);
        std::unique_ptr<std::ostream> headerStream = getOstreamForFileName(lexerName + ".h");
        std::unique_ptr<std::ostream> implementationStream  = getOstreamForFileName(lexerName + ".cpp");

        std::map<const std::string, templ::TemplateContext> topLevel;
        
        topLevel["name"] = templ::make_string(lexerName);

        //The DEADSTATE gets a brand new state: dfa.numStates
        topLevel["reject_state"] = make_State(dfa.numStates);
        topLevel["num_states"] = make_State(dfa.numStates + 1);

        auto transition_indices = buildTransitionIndices(dfa);
        topLevel["trans_idx"] = transformTransitionIndices(transition_indices.first);
        topLevel["num_transitions_per_state"] = make_int(transition_indices.second);

        topLevel["table"] = buildTable(dfa, transition_indices.first, transition_indices.second);

        topLevel["token_types"] = buildTokenList(dfa);

        std::vector<templ::TemplateContext> acTokens;
        for (State s = 0; s < dfa.numStates; s++) {
            if (dfa.accepting.count(s)) {
                acTokens.push_back(make_map_elem("token", dfa.acceptingToken.find(s)->second));
            } else {
                acTokens.push_back(make_map_elem("token", "nonmatching"));
            }
        }
        acTokens.push_back(make_map_elem("token", "nonmatching"));
        topLevel["tokens"] = templ::make_array(acTokens);

        templ::TemplateContext topLevelMap = templ::make_map(topLevel);

        doTemplate(*headerStream, "c++/lexer.h", topLevelMap);
        doTemplate(*implementationStream, "c++/lexer.cpp", topLevelMap);
    }

    templ::TemplateContext CppBackend::buildTable(const DFA& dfa, const std::vector<unsigned char>& transition_idx, int num_transitions_per_state) const {
        std::map<unsigned char, unsigned char> reverse_trans;
        for (int i = 0; i < 256; i++) {
            reverse_trans[transition_idx[i]] = i;
        }

        std::vector<templ::TemplateContext> table;
        
        for (State s = 0; s < dfa.numStates; s++) {
            std::vector<templ::TemplateContext> row;
            for (int i = 0; i < num_transitions_per_state; i++) {
                State to = dfa.delta.find(s)->second.find(reverse_trans[i])->second;
                if (to == deadState)
                    to = dfa.numStates; //The new Dead state name
                row.push_back(templ::make_map({{"state", make_State(to)}}));
            }
            table.push_back(templ::make_map({{"row", templ::make_array(row)}}));
        }
        // NOTE: there is no transition table entry for the dead state
        //       since the algorithm should never loop in the dead state

        return templ::make_array(table);
    }

    templ::TemplateContext CppBackend::buildTokenList(const DFA& dfa) const {
        std::set<std::string> tokens;
        for (const auto& pr : dfa.acceptingToken) {
            tokens.insert(pr.second);
        }
        tokens.insert("ignore");

        std::vector<templ::TemplateContext> tokenList;
        for (const std::string& s : tokens) {
            tokenList.push_back(make_map_elem("type", s));
        }
        return templ::make_array(tokenList);
    }

    std::pair<std::vector<unsigned char>, int> CppBackend::buildTransitionIndices(const DFA& dfa) const {
        std::vector<unsigned char> transition_idx(256);
        int count = 0;
        for (int i = 0; i < 256; i++) {
            int found = i;
            for (int j = 0; j < i; j++) {
                bool differ = false;
                for (State s = 0; s < dfa.numStates; s++) {
                    if (dfa.delta.find(s)->second.find(i)->second != dfa.delta.find(s)->second.find(j)->second) {
                        differ = true;
                        break;
                    }
                }
                if (!differ) {
                    found = j;
                    break;
                }
            }
            if (found != i)
            {
                transition_idx[i] = transition_idx[found];
            } else {
                transition_idx[i] = count++;
            }
        }

        return std::make_pair(transition_idx, count);
    }

    templ::TemplateContext CppBackend::transformTransitionIndices(const std::vector<unsigned char>& transition_indices) const {
        std::vector<templ::TemplateContext> new_trans;
        for (auto& i : transition_indices) {
            new_trans.push_back(make_map_elem("trans", "(unsigned char)" + std::to_string(i)));
        }
        return templ::make_array(new_trans);
    }

} } //namespace lxs::backends
