#include <cctype>
#include <cstdlib>
#include <gelf.h>
#include <err.h>
#include <sysexits.h>
#include <fcntl.h>
#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>
#include <getopt.h>
#include <inttypes.h>
#include <sstream>
#include <memory>
#include <map>
#include <elf.h>

#include "trygvis/elfinfo/Ld.h"

using namespace std;
using namespace trygvis::elfinfo;

struct Usage {
    uint64_t maxAddress;
};

std::vector<MemoryArea> memoryAreas;
std::map<string, shared_ptr<Usage>> usages;

char *filename = nullptr;

char *ld_filename = nullptr;

char *program;

__attribute__((noreturn))
void usage(const char *reason = nullptr) {
    if (reason) {
        fprintf(stderr, "%s\n", reason);
    }
    fprintf(stderr, "usage: %s -f file [-l ld] [-t start:size] [-d start:size]\n", program);
    fprintf(stderr, "  -t/-d/-b: add text/data area\n");
    fprintf(stderr, "At least one area or a ld file has to be specified\n");
    exit(EX_USAGE);
}

void parse_start_size(char *input, uint64_t &start, uint64_t &size) {
    char *str_size = strchr(input, ':');

    if (!str_size) {
        usage("bad area specification, missing ':'");
    }

    *str_size = '\0';
    str_size++;

    if (sscanf(input, "%" SCNi64, &start) != 1) {
        usage("bad area specification, could not parse start number");
    }

    size_t str_size_len = strlen(str_size);

    if (str_size_len < 1) {
        usage("bad area specification");
    }

    char suffix = str_size[str_size_len - 1];
    int modifier;

    if (!isdigit(suffix)) {
        switch (suffix) {
            case 'k':
            case 'K':
                modifier = 1024;
                break;
            case 'm':
            case 'M':
                modifier = 1024 * 1024;
                break;
            case 'g':
            case 'G':
                modifier = 1024 * 1024 * 1024;
                break;
            default:
                usage("bad size modifier, only 'k', 'M' and 'G' are allowed");
        }
    } else {
        modifier = 1;
    }

    if (sscanf(str_size, "%" SCNu64, &size) != 1) {
        usage("bad area specification, could not parse size number");
    }
    size = size * modifier;
}

bool debug = false;

void parse_args(int argc, char **argv) {
    int c;
    static int anonymousTextCounter = 0, anonymousDataCounter = 0;

    while ((c = getopt(argc, argv, "Df:l:t:d:")) != -1) {
        switch (c) {
            case 'D':
                debug = true;
                break;
            case 't':
            case 'd': {
                Elf64_Addr start;
                Elf64_Xword size;
                parse_start_size(optarg, start, size);

                string name;
                set<MemoryAttribute> attributes;
                if (c == 't') {
                    attributes.insert(MemoryAttribute::R);
                    attributes.insert(MemoryAttribute::W);
                    attributes.insert(MemoryAttribute::X);

                    name = "text";
                    if (anonymousTextCounter > 0) {
                        name += anonymousTextCounter;
                    }
                    anonymousTextCounter++;
                } else {
                    attributes.insert(MemoryAttribute::R);
                    attributes.insert(MemoryAttribute::W);

                    name = "data";
                    if (anonymousDataCounter > 0) {
                        name += anonymousDataCounter;
                    }
                    anonymousDataCounter++;
                }
                memoryAreas.push_back(MemoryArea {
                    .name = name,
                    .origin = start,
                    .length = size,
                    .attributes = attributes
                });
                break;
            }
            case 'f':
                filename = optarg;
                break;
            case 'l':
                ld_filename = optarg;
                break;
            case '?':
                if (optopt == 'c')
                    errx(EX_USAGE, "Option -%c requires an argument.\n", optopt);
                else
                    errx(EX_USAGE, "Unknown option `-%c'.\n", optopt);
            default:
                abort();
        }
    }

    if (!filename || (memoryAreas.empty() && !ld_filename)) {
        usage();
    }
}

string to_iso(uint64_t i) {
    const char *suffix;
    std::stringstream buf;
    if (i >= 1024 * 1024 * 1024) {
        i /= 1024 * 1024 * 1024;
        suffix = "G";
    } else if (i >= 1024 * 1024) {
        i /= 1024 * 1024;
        suffix = "M";
    } else if (i >= 1024) {
        i /= 1024;
        suffix = "k";
    } else {
        suffix = "";
    }
    buf << i << suffix;
    return buf.str();
}

int main(int argc, char **argv) {
    program = argv[0];
    parse_args(argc, argv);

    if (ld_filename) {
        try {
            LdScriptLoader loader;
            loader.setDebug(debug);
            LdScript file = loader.load(ld_filename);

            memoryAreas.insert(memoryAreas.end(), file.memoryAreas.begin(), file.memoryAreas.end());
        } catch (LdParseException &e) {
            for (auto &m: e.messages) {
                fprintf(stderr, "%s\n", m.c_str());
            }
            errx(EX_DATAERR, "Could not parse ld script");
        } catch (std::exception &e) {
            errx(EX_DATAERR, "Unhandled exception when parsing LD script. "
                "The parser is new and fragile, so please file a bug and send me your LD script: %s", e.what());
        } catch (...) {
            errx(EX_DATAERR, "Unhandled exception when parsing LD script. "
                "The parser is new and fragile, so please file a bug and send me your LD script.");
        }
    }

    if (elf_version(EV_CURRENT) == EV_NONE)
        errx(EX_SOFTWARE, "ELF library initialization failed: %s", elf_errmsg(-1));

    int fd;
    if ((fd = open(filename, O_RDONLY, 0)) < 0)
        err(EX_NOINPUT, "open \"%s\" failed", argv[1]);

    Elf *e;
    if ((e = elf_begin(fd, ELF_C_READ, NULL)) == NULL)
        errx(EX_SOFTWARE, "elf_begin() failed: %s.", elf_errmsg(-1));
    if (elf_kind(e) != ELF_K_ELF)
        errx(EX_DATAERR, "%s is not an ELF object.", argv[1]);

    size_t shstrndx;
    if (elf_getshdrstrndx(e, &shstrndx) != 0)
        errx(EX_SOFTWARE, "elf_getshdrstrndx() failed: %s.", elf_errmsg(-1));

    size_t program_header_count;
    if (elf_getphdrnum(e, &program_header_count) != 0)
        errx(EX_DATAERR, "elf_getphdrnum() failed: %s.", elf_errmsg(-1));

    for (int i = 0; i < program_header_count; i++) {
        GElf_Phdr phdr;

        if (gelf_getphdr(e, i, &phdr) != &phdr)
            errx(EX_SOFTWARE, "getphdr() failed: %s.", elf_errmsg(-1));

        if (phdr.p_type == PT_LOAD) {
            if (phdr.p_flags == (PF_X | PF_W | PF_R) || phdr.p_flags == (PF_X | PF_R)) {
                if (debug) {
                    printf("Adding PH #%d as text\n", i);
                }

            } else if (phdr.p_flags == (PF_R | PF_W)) {
                if (phdr.p_filesz > 0) {
                    if (debug) {
                        printf("Adding PH #%d as data\n", i);
                    }
                }
                else {
                    if (debug) {
                        printf("Adding PH #%d as bss\n", i);
                    }
                }
            } else {
                warnx("Unknown flag combination: 0x%02x", phdr.p_flags);
                continue;
            }

            auto ma = std::find_if(memoryAreas.begin(), memoryAreas.end(), [&](MemoryArea &memoryArea) {
                return memoryArea.contains(phdr.p_vaddr);
            });

            if (ma == memoryAreas.end()) {
                warnx(
                    "Could not find a area for elf header #%d of type %s, at address 0x%08" PRIx64 " with size %" PRId64 "\n",
                    i, ""/*to_str(expectedType)*/, phdr.p_vaddr, phdr.p_memsz);
            } else {
                if (debug && !ma->contains(phdr.p_vaddr + phdr.p_memsz)) {
                    warnx("Area overflow: %s", ma->name.c_str());
                }

                auto &usage = usages[ma->name];

                if (!usage) {
                    usage = usages[ma->name] = make_shared<Usage>();
                }

                usage->maxAddress = phdr.p_memsz;
            }
        } else {
            // ignored
            if (debug) {
                printf("Ignoring ELF section #%d\n", i);
            }
        }
    }

    printf("Memory areas\n");
    printf("Name       Flags Start    End       Size Used\n");
    std::for_each(memoryAreas.begin(), memoryAreas.end(), [&](const MemoryArea &s) {
        auto &usage = usages[s.name];

        int used_pct = 0;
        if (usage) {
            used_pct = (int) (double(usage->maxAddress) / double(s.length) * 100.0);
        }
        printf("%-10s %s   %08" PRIx64 " %08" PRIx64 " %5s %3d%%\n",
               s.name.c_str(),
               s.attributes_string().c_str(),
               s.origin,
               s.origin + s.length,
               to_iso(s.length).c_str(),
               used_pct);
    });

//    printf("\n");
//    printf("Size by type\n");
//    printf("text=%zu, data=%zu, bss=%zu\n", text_size, data_size, bss_size);
    return EXIT_SUCCESS;
}