#include "molasses/lexer.h"
#include "molasses/parser_primitives.h"
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <stack>
#include <string>
#include <variant>

using compile_element = std::variant<molasses::lexed_output, molasses::generate_context, std::string>;

int main(int argc, char** argv) {
	std::vector<std::string> arguments;
    while(argc > 1) {
        argv++;
        arguments.emplace_back(*argv, strlen(*argv));
        argc--;
    }

    try {

        std::stack<compile_element> compile_stack;

        for(auto elem : arguments) {
            if(elem == "generate") {
                if(std::holds_alternative<std::string>(compile_stack.top())) {
                    auto filename = std::get<std::string>(compile_stack.top());
                    compile_stack.pop();
                    if(std::holds_alternative<molasses::generate_context>(compile_stack.top())) {
                        auto generator = std::get<molasses::generate_context>(compile_stack.top());
                        compile_stack.pop();
                        auto assembler = molasses::generate(generator);
                        std::ofstream output(filename + ".s");
                        for(const auto& line : assembler) {
                            output << line;
                        }
                        compile_stack.emplace(filename);
                    } else
                        throw std::runtime_error("generate expects a parsed output");
                } else
                    throw std::runtime_error("generate expects a filename");
            } else if(elem == "parse") {
                molasses::parser_context ctx;
                ctx = molasses::register_integers(ctx);
                ctx = molasses::register_i32_operations(ctx);

                if(std::holds_alternative<molasses::lexed_output>(compile_stack.top())) {
                    auto lexer = std::get<molasses::lexed_output>(compile_stack.top());
                    compile_stack.pop();
                    auto generator = molasses::parse(ctx, lexer);
                    compile_stack.emplace(generator);
                } else
                    throw std::runtime_error("parse expects a lexed output");
            } else if(elem == "lex") {
                if(std::holds_alternative<std::string>(compile_stack.top())) {
                    auto filename = std::get<std::string>(compile_stack.top());
                    compile_stack.pop();
                    if(not std::filesystem::exists(filename))
                        throw std::runtime_error("file " + filename + " does not exist");
                    std::ifstream t(filename);
                    std::stringstream buffer;
                    buffer << t.rdbuf();
                    auto lexed = molasses::lex(filename, buffer.str());
                    compile_stack.emplace(lexed);
                } else
                    throw std::runtime_error("lex expects a filename");
            } else if(elem == "lex-all") {
                std::vector<molasses::lexed_output> lexed_list;
                while(not compile_stack.empty() and std::holds_alternative<std::string>(compile_stack.top())) {
                    auto filename = std::get<std::string>(compile_stack.top());
                    compile_stack.pop();
                    if(not std::filesystem::exists(filename))
                        throw std::runtime_error("file " + filename + " does not exist");
                    std::ifstream t(filename);
                    std::stringstream buffer;
                    buffer << t.rdbuf();
                    auto lexed = molasses::lex(filename, buffer.str());
                    lexed_list.emplace_back(lexed);
                }
                for(auto& lexed : lexed_list) {
                    compile_stack.emplace(std::move(lexed));
                }
            } else if(elem == "merge") {
                if(std::holds_alternative<molasses::lexed_output>(compile_stack.top())) {
                    auto lexer_1 = std::get<molasses::lexed_output>(compile_stack.top());
                    compile_stack.pop();
                    if(std::holds_alternative<molasses::lexed_output>(compile_stack.top())) {
                        auto lexer_2 = std::get<molasses::lexed_output>(compile_stack.top());
                        compile_stack.pop();
                        compile_stack.emplace(molasses::concatenate(lexer_1, lexer_2));
                    } else
                        throw std::runtime_error("merge expects 2 lexed outputs");
                } else
                    throw std::runtime_error("merge expects 2 lexed outputs");
            } else if(elem == "merge-all") {
                if(std::holds_alternative<molasses::lexed_output>(compile_stack.top())) {
                    auto lexer_1 = std::get<molasses::lexed_output>(compile_stack.top());
                    compile_stack.pop();
                    while(not compile_stack.empty() and std::holds_alternative<molasses::lexed_output>(compile_stack.top())) {
                        auto lexer_2 = std::get<molasses::lexed_output>(compile_stack.top());
                        compile_stack.pop();
                        lexer_1 = molasses::concatenate(lexer_1, lexer_2);
                    }
                    compile_stack.emplace(lexer_1);
                } else
                    throw std::runtime_error("merge-all expects at least 1 lexed outputs");
            } else if(elem == "assemble") {
                if(std::holds_alternative<std::string>(compile_stack.top())) {
                    auto filename = std::get<std::string>(compile_stack.top());
                    compile_stack.pop();
                    std::stringstream compile;
                    compile << "clang -c " << filename << ".s -o " << filename << ".o";
                    std::stringstream link;
                    link << "ld -e _start " << filename << ".o -o " << filename;
                    std::cout << compile.str() << std::endl;
                    system(compile.str().c_str());
                    std::cout << link.str() << std::endl;
                    system(link.str().c_str());
                } else
                    throw std::runtime_error("assemble expects an assembly file");
            } else if(elem == "help" or elem == "--help") {
                std::cout << "# Sugar\n\n";
                std::cout << "## Commands\n\n";
                std::cout << "lex : string ➔ lexed_output\n";
                std::cout << "> takes a filename to a file that must be compiled\n\n";
                std::cout << "lex-all : string* ➔ lexed_output*\n";
                std::cout << "> takes as many filenames to files that must be compiled as can be read and passes them through the lexed\n\n";
                std::cout << "merge : lexed_output lexed_output ➔ lexed_output\n";
                std::cout << "> merges two lexed modules together\n\n";
                std::cout << "merge-all : lexed_output* ➔ lexed_output\n";
                std::cout << "> merges as many lexed modules together as present on the top of the stack\n\n";
                std::cout << "parse : lexed_output ➔ parsed_output\n";
                std::cout << "> prepares code for generation\n\n";
                std::cout << "generate : parsed_output string ➔ string\n";
                std::cout << "> takes a root filename, it will be appended with \".s\" and that will be the generated assembly file,\n";
                std::cout << "> the filename will not be consumed\n\n";
                std::cout << "assemble : string ➔ _ \n";
                std::cout
                    << "> takes a root filename, it will be appended with \".s\" and that file will be compiled,\n";
                std::cout << "> the compiled output will be the given filename\n\n";
                std::cout << "help : _ ➔ _ \n";
                std::cout << "> prints this help\n\n";
                std::cout << "## Examples\n\n";
                std::cout << "- compile the file \"example.mol\" into the \"potato.s\" assembly file\n";
                std::cout << "> `$ sugar example.mol lex parse potato generate`\n";
                std::cout << "\n";
                std::cout << "- compile the file \"example.mol\" into the \"potato\" executable\n";
                std::cout << "> `$ sugar example.mol lex parse potato generate assemble`\n";
                std::cout << "\n";
                std::cout << "- compile the file \"example.mol\" and \"2.mol\" into the \"potato\" executable\n";
                std::cout << "> `$ sugar example.mol lex 2.mol lex merge parse potato generate assemble`\n";
            } else
                compile_stack.emplace(elem);
        }

        if(compile_stack.size() > 1) throw std::runtime_error("build left unfinished operations");
        if(not compile_stack.empty()) {
            if(std::holds_alternative<molasses::lexed_output>(compile_stack.top())) {
                auto lexer = std::get<molasses::lexed_output>(compile_stack.top());
                for(auto elem : lexer.symbols) {
                    std::cout << elem << " ";
                }
                std::cout << "\n\n";
                for(auto [id, name] : lexer.dictionary) {
                    std::cout << id << ": " << name << "\n";
                }
            } else if(std::holds_alternative<molasses::generate_context>(compile_stack.top())) {
                auto generator = std::get<molasses::generate_context>(compile_stack.top());
                for(const auto& elem : generator.procedures) {
                    std::cout << elem->_name << " : ";
                    for(const auto& args : elem->_args) {
                        std::cout << args << " ";
                    }
                    std::cout << "->";
                    for(const auto& rets : elem->_rets) {
                        std::cout << " " << rets;
                    }
                    std::cout << "\n";
                }
            }
        }
    } catch (molasses::parser_error& error) {
        std::cerr << error.what();
    }
}