#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <regex>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <ctype.h>
using namespace std;

int main(int argc, char ** argv)
{
    std::vector<const char*> gccargs;
    const char * input  = nullptr;
    const char * output = nullptr;
    const char * gcc    = "gcc";
    bool         m32    = false;

    /*parse command line arguments*/ {
        for(int i = 1; i < argc; ++i) {
            const char * A = argv[i];
            if (A[0] == '-') {
                if (A[1] == 'c' && A[2] == 'c' && !A[3]) {
                    if (++i == argc) {
                        std::cerr<<"Error: compiler must be specified after -cc"<<std::endl;
                        return 1;
                    }
                    gcc = argv[i];
                    continue;
                }
                if (!strcmp(A, "-x32")) {
                    m32 = true;
                    continue;
                }
                if (A[1] == 'o' && !A[2]) {
                    if (output) {
                        std::cerr<<"Error: more than one output file specified"<<std::endl;
                        return 1;
                    }
                    if (++i == argc) {
                        std::cerr<<"Error: output file must be specified after -o"<<std::endl;
                        return 1;
                    }
                    output = argv[i];
                    continue;
                }
                gccargs.push_back(A);
                continue;
            }
            if (input) {
                std::cerr<<"Error: more than one input file specified"<<std::endl;
                return 1;
            }
            input = A;
        }
    }
    if (!input || !output) {
        std::cerr<<"Error: input and output files must be specified"<<std::endl;
        return 1;
    }

    /*read beginning of input file for #pragma CODE32*/ {
        std::ifstream input_file(input);
        if (!input_file.is_open()) {
            std::cout<<"Error: failed to open input file"<<std::endl;
            return 1;
        }
        for(std::string line; std::getline(input_file, line); ) {
            if (line.empty() || line[0] != '#') break;
            if (line == "#define MULTIBOOT_INTERFACE") {
                std::cout<<"file is compiled in 32bit mode"<<std::endl;
                m32 = true;
                break;
            }
        }
    }

    std::string tmp(output);
    tmp.append("_in.s");
    std::string tmp2(output);
    tmp2.append("_out.s");

    /*run compiler to generate assembly code*/ {
        int pid = fork();
        if (pid < 0) {
            auto c = strerror(errno);
            std::cerr<<"Error: failed to create new process: "<<c<<std::endl;
            return 1;
        }
        if (!pid) {
            gccargs.insert(gccargs.begin(), gcc);
            gccargs.push_back("-fPIC");
            gccargs.push_back("-fno-rtti");
            gccargs.push_back("-fno-exceptions");
            gccargs.push_back("-fno-asynchronous-unwind-tables");
            gccargs.push_back("-ffreestanding");
            if (m32) gccargs.push_back("-m32");
            gccargs.push_back("-S");
            gccargs.push_back("-o");
            gccargs.push_back(tmp.c_str());
            gccargs.push_back(input);
            gccargs.push_back(nullptr);

            int ret = execvp(gcc,(char * const * )gccargs.data());
            if (ret < 0) {
                auto c = strerror(errno);
                std::cerr<<"Error: failed to create new process: "<<c<<std::endl;
                return 1;
            }
        } else {
            int status;
            if (pid != waitpid(pid, &status, 0)) {
                auto c = strerror(errno);
                std::cerr<<"Error: failed to wait for compiler: "<<c<<std::endl;
                kill(pid, SIGKILL);
                return 1;
            }
            if (!WIFEXITED(status) || WEXITSTATUS(status)) {
                std::cerr<<"Error: compiler terminated unsuccessfully"<<std::endl;
                return 1;
            }
        }
    }
    std::ifstream tmp_in(tmp.c_str());
    std::ofstream tmp_out(tmp2.c_str());
    if (!tmp_in.is_open() || !tmp_out.is_open()) {
        std::cerr<<"Error: failed to open temporary file"<<std::endl;
        return 1;
    }

    /*generate fake GOT*/ {
        tmp_out<<"\t.text"<<std::endl;
        tmp_out<<".Ltext_offset:"<<std::endl<<std::endl;
    }

    /*modify sources to become GOTless*/ {
        std::regex is_stuff("^\\s*(\\#|\\.file).*");
        std::regex is_text("^\\s*\\.text.*");
        std::regex is_add_got("^\\s*addl\\s+\\$_GLOBAL_OFFSET_TABLE_\\,\\s*(\\%[a-z0-9]+)\\s*.*$");
        std::regex is_ref_got("^\\s*leal\\s+([\\._]{0,1}[A-Za-z0-9_]+)\\@GOTOFF\\((\\%[a-z0-9]+)\\)\\,\\s*(\\%[a-z0-9]+)\\s*$");
        std::regex is_unexp(".*(_GLOBAL_OFFSET_TABLE_|GOTOFF).*");

        int linenumber = 0;
        bool fail = false;
        for(std::string line; std::getline(tmp_in, line); ) {
            ++linenumber;
            //test for #xxxxx
            if (std::regex_match(line, is_stuff)) {
                continue;
            }
            if (m32 && std::regex_match(line, is_text)) {
                tmp_out<<"\t.text"<<std::endl<<"\t.code32"<<std::endl;
                continue;
            }
            std::smatch mr;
            //find "addl	$_GLOBAL_OFFSET_TABLE_, %reg"
            if (std::regex_match(line, mr, is_add_got)) {
                if (mr.size() == 2) {
                    tmp_out<<"//"<<line<<std::endl;
                    tmp_out<<"\taddl\t$(.Ltext_offset - .), "<<mr[1].str()<<std::endl;
                    continue;
                } else {
                    tmp_out<<"//ERROR!"<<std::endl;
                    std::cerr<<"Error["<<tmp<<":"<<linenumber<<"]: unsupported instruction"<<std::endl;
                    fail = true;
                }
                continue;
            }
            //find usage of GOT reference and replace it with direct link
            else if (std::regex_match(line, mr, is_ref_got)) {
                if (mr.size() == 4) {
                    //leal does not work here, AS cannot .set XXX, label-.Ltext_offset
                    //therefore we cannot use it in leal XXX(%reg), %dst
                    tmp_out<<"//"<<line<<std::endl;
                    tmp_out<<"\tmovl\t$("<<mr[1].str()<<"-.Ltext_offset), "<<mr[3].str()<<std::endl;
                    tmp_out<<"\taddl\t"<<mr[2].str()<<", "<<mr[3].str()<<std::endl;
                    continue;
                } else {
                    tmp_out<<"//ERROR!"<<std::endl;
                    std::cerr<<"Error["<<tmp<<":"<<linenumber<<"]: unsupported instruction"<<std::endl;
                    fail = true;
                }
            }
            else if (std::regex_match(line, is_unexp)) {
                tmp_out<<"//ERROR!"<<std::endl;
                std::cerr<<"Error["<<tmp<<":"<<linenumber<<"]: unsupported instruction"<<std::endl;
                fail = true;
            }

            tmp_out<<line<<std::endl;
        }
        if (fail)
            return 1;
    }

    /*run compiler to generate object code*/ {
        int pid = fork();
        if (pid < 0) {
            auto c = strerror(errno);
            std::cerr<<"Error: failed to create new process: "<<c<<std::endl;
            return 1;
        }
        if (!pid) {
            std::vector<const char*> args;
            args.push_back(gcc);
            args.push_back("-c");
            args.push_back("-o");
            args.push_back(output);
            args.push_back(tmp2.c_str());
            args.push_back(nullptr);

            int ret = execvp(gcc,(char * const * )args.data());
            if (ret < 0) {
                auto c = strerror(errno);
                std::cerr<<"Error: failed to create new process: "<<c<<std::endl;
                return 1;
            }
        } else {
            int status;
            if (pid != waitpid(pid, &status, 0)) {
                auto c = strerror(errno);
                std::cerr<<"Error: failed to wait for compiler: "<<c<<std::endl;
                kill(pid, SIGKILL);
                return 1;
            }
            if (!WIFEXITED(status) || WEXITSTATUS(status)) {
                std::cerr<<"Error: compiler terminated unsuccessfully"<<std::endl;
                return 1;
            }
        }
    }
    return 0;
}