#include <iostream>
#include <vector>
#include <cstring>
#include <algorithm>
#include "Options.h"
#include "Narc.h"
#include "RelocElfReader.h"

Options::Options(int argc, char **argv) {
    for (int i = 1; i < argc; i++) {
        string arg{argv[i]};
        if (arg == "-f" || arg == "--flatten") {
            flatten = true;
        } else if (arg == "-p" || arg == "--padding") {
            int padval_i = stoi(argv[++i]);
            if (padval_i < 0 || padval_i > 255) {
                throw command_error(string{"invalid 8-bit value "} + argv[i] + " for " + arg);
            }
            padval = static_cast<char>(padval_i);
        } else if (arg == "-n" || arg == "--naix") {
            naix = true;
        } else if (arg[0] == '-') {
            throw command_error("unrecognized option flag: " + arg);
        } else if (posargs.size() >= 2) {
            throw command_error("unrecognized positional argument: " + arg);
        } else {
            posargs.emplace_back(arg);
        }
    }
    if (posargs.size() < 2) {
        throw command_error("missing positional arg");
    }
    objfile.open(posargs[0], ios::in | ios::binary);
    narcfile.open(posargs[1], ios::out | ios::binary);
}

void Options::ReadObjectFile(vector<unsigned char> &rodata, vector<uint32_t> &sizes) {
    ELF_ASSERT(objfile.HasSection(".rodata"));
    rodata.resize(objfile.GetSectionHeader(".rodata").sh_size);
    objfile.ReadSectionData(objfile.GetSectionHeader(".rodata"), rodata.data());
    // Determine what O file we're dealing with
    if (objfile.HasSymbol("__size")) {
        uint32_t size = objfile.GetSymbol("__size").st_size;
        if (size == sizeof(uint32_t)) {
            objfile.ReadSymbolData(objfile.GetSymbol("__size"), &size);
            sizes.resize((rodata.size() + size - 1) / size);
            fill(sizes.begin(), sizes.end(), size);
        } else {
            sizes.resize(objfile.GetSymbol("__size").st_size / sizeof(uint32_t));
            objfile.ReadSymbolData(objfile.GetSymbol("__size"), sizes.data());
        }
    } else {
        auto pred = [&](const Elf32_Sym &sym) {
            return sym.st_size != 0
                   && strcmp(objfile.GetSectionName(objfile.sections()[sym.st_shndx]), ".rodata") == 0
                   && strcmp(objfile.GetSymbolName(sym), "__size") != 0
                   && strcmp(objfile.GetSymbolName(sym), "__data") != 0
                   && strcmp(objfile.GetSymbolName(sym), ".rodata") != 0;
        };
        sizes.resize(count_if(objfile.symbols().begin(), objfile.symbols().end(), pred));
        ELF_ASSERT(!sizes.empty());
        int t = 0;
        for (const auto &sym : objfile.symbols()) {
            if (pred(sym)) {
                sizes[t++] = sym.st_size;
            }
        }
    }
}

void Options::OverwritePadding(vector<unsigned char> &rodata, vector<uint32_t> &sizes) const {
    uint32_t end = 0;
    for (auto & size : sizes) {
        end += size;
        uint32_t pad_end = (end + 3) & ~3;
        memset(&rodata[end], padval, pad_end - end);
        end = pad_end;
    }
}

void Options::WriteNarc(vector<unsigned char> &rodata, vector<uint32_t> &sizes) {
    if (!flatten) {
        FileImages fimg(rodata);
        FileNameTableEntry fntent;
        FileNameTable fnt;
        vector<FileAllocationTableEntry> fatent = FileAllocationTableEntry::_make(sizes);
        FileAllocationTable fat(fatent);
        NarcHeader narc(fat, fnt, fimg);
        narcfile.write((char *)&narc, sizeof(narc));
        narcfile.write((char *)&fat, sizeof(fat));
        narcfile.write((char *)fatent.data(), fatent.size() * sizeof(FileAllocationTableEntry));
        narcfile.write((char *)&fnt, sizeof(fnt));
        narcfile.write((char *)&fntent, sizeof(fntent));
        narcfile.write((char *)&fimg, sizeof(fimg));
    }
    narcfile.write((char *)rodata.data(), rodata.size());
}

void Options::WriteNaix(vector<uint32_t> &sizes) {
    if (naix) {
        string naixname = posargs[1].substr(0, posargs[1].find_last_of('.')) + ".naix";
        string stem = naixname.substr(naixname.find_last_of('/') + 1, naixname.find_last_of('.') - naixname.find_last_of('/') - 1);
        string stem_upper = stem;
        for (auto &c : stem_upper) { c = toupper(c); }
        ofstream naixfile(naixname);
        naixfile << "/*\n"
                    " * THIS FILE WAS AUTOMATICALLY\n"
                    " *  GENERATED BY tools/o2narc\n"
                    " *      DO NOT MODIFY!!!\n"
                    " */\n"
                    "\n"
                    "#ifndef NARC_" << stem_upper << "_NAIX_\n"
                                                     "#define NARC_" << stem_upper << "_NAIX_\n"
                                                                                      "\n"
                                                                                      "enum {\n";
        char num_buf[9] = "00000000";
        for (int i = 0; i < sizes.size(); i++) {
            naixfile << "    NARC_" << stem << "_" << stem << "_" << num_buf << " = " << i << "," << endl;
            for (int k = 7; k >= 0; k--) {
                num_buf[k]++;
                if (num_buf[k] > '9') {
                    num_buf[k] = '0';
                } else {
                    break;
                }
            }
        }
        naixfile << "};\n\n#endif //NARC_" << stem_upper << "_NAIX_\n";
    }
}

int Options::main() {
    vector<uint32_t> sizes;
    vector<unsigned char> rodata;

    ReadObjectFile(rodata, sizes);
    OverwritePadding(rodata, sizes);
    WriteNarc(rodata, sizes);
    WriteNaix(sizes);
    return 0;
}