#include <iostream>
#include <fstream>
#include <algorithm>
#include <cstdlib>
#include <vector>

static void ltrim(std::string &s) {
    s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](char c) {
        return !std::isspace(c);
    }));
}

static void rtrim(std::string &s) {
    s.erase(std::find_if(s.rbegin(), s.rend(), [](char c) {
        return !std::isspace(c);
    }).base(), s.end());
}

static void trim(std::string &s) {
    ltrim(s);
    rtrim(s);
}

static bool startsWith(const std::string &s, const std::string &start) {
    return s.rfind(start, 0) == 0;
}

static bool endsWith(const std::string &s, const std::string &end) {
    if (s.length() >= end.length()) {
        return s.compare(s.length() - end.length(), end.length(), end) == 0;
    } else {
        return false;
    }
}

static bool extractNumber(std::string macroValue, uint64_t &output) {
    while (startsWith(macroValue, "(") && endsWith(macroValue, ")")) {
        macroValue = macroValue.substr(1, macroValue.length() - 2);
    }

    if (startsWith(macroValue, "(NTSTATUS)")) {
        macroValue = macroValue.substr(10);
        ltrim(macroValue);
    }

    errno = 0;
    auto value = std::strtoull(&macroValue[0], nullptr, 0);
    if (errno != 0) {
        return false;
    }

    output = value;
    return true;
}

struct NtStatusCode {
    explicit NtStatusCode() = default;

    explicit NtStatusCode(std::string name, uint64_t value) : name(std::move(name)), value(value) {}

    std::string name;
    uint64_t value = 0;
};

static std::vector<NtStatusCode> predefinedCodes() {
    // Some codes are not in ntstatus.h for some reason...
    // https://docs.microsoft.com/en-us/windows-hardware/drivers/debugger/specific-exceptions
    return {
            NtStatusCode{"STATUS_APPLICATION_HANG", 0xCFFFFFFF},
            NtStatusCode{"STATUS_CPP_EH_EXCEPTION", 0xE06D7363},
            NtStatusCode{"STATUS_CLR_EXCEPTION", 0xE0434f4D},
    };
}

int main(int argc, const char **argv) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <header-out> <source-out>" << std::endl;
        return 1;
    }

    std::ifstream ntstatusTarget(NTSTATUS_PREPROCESSOR_OUT);
    if (!ntstatusTarget.is_open()) {
        std::cerr << "Failed to open preprocessor output at " << NTSTATUS_PREPROCESSOR_OUT << std::endl;
        return 1;
    }

    std::string ntstatusPath;

    std::string line;
    while (std::getline(ntstatusTarget, line)) {
        trim(line);

        if (endsWith(line, "ntstatus.h")) {
            ntstatusPath = line;
            break;
        }
    }

    ntstatusTarget.close();

    if (ntstatusPath.empty()) {
        std::cerr << "Failed to find path to ntstatus.h in generated preprocessor output" << std::endl;
        return 1;
    }


    std::cout << "nstatus.h at " << ntstatusPath << std::endl;
    std::ifstream ntstatusHeader(ntstatusPath);
    if (!ntstatusHeader.is_open()) {
        std::cerr << "Failed to open ntstatus.h" << std::endl;
        return 1;
    }

    std::vector<NtStatusCode> codes = predefinedCodes();

    while (std::getline(ntstatusHeader, line)) {
        trim(line);

        if (startsWith(line, "#define") && line.find("NTSTATUS") != std::string::npos) {
            line = line.substr(7);
            ltrim(line);

            auto space = line.find(' ');
            if (space == std::string::npos) {
                std::cerr << "Skipping #define " << line << " as the macro has no value" << std::endl;
                continue;
            }

            auto name = line.substr(0, space);
            auto value = line.substr(space + 1);
            ltrim(value);

            NtStatusCode code;
            code.name = name;

            if (!extractNumber(value, code.value)) {
                std::cerr << "Skipping #define " << line << " because its value couldn't be parsed" << std::endl;
            } else {
                codes.emplace_back(std::move(code));
            }
        }
    }

    std::cout << "Found " << codes.size() << " NTSTATUS codes" << std::endl;

    std::ofstream outputHeader(argv[1]);
    if (!outputHeader.is_open()) {
        std::cerr << "Failed to open header output file " << argv[1] << std::endl;
        return 1;
    }

    outputHeader << "// AUTO GENERATED FILE, DO NOT EDIT!" << std::endl;
    outputHeader << "// This file has been generated by nstatus-gen from the systeminfo library" << std::endl;
    outputHeader << std::endl;
    outputHeader << "#pragma once" << std::endl;
    outputHeader << std::endl;
    outputHeader << "#include <cstdint>" << std::endl;
    outputHeader << "#include <string>" << std::endl;
    outputHeader << std::endl;
    outputHeader << "namespace Sys {" << std::endl;
    outputHeader << "namespace Win32 {" << std::endl;
    outputHeader << "bool lookupNtStatusCodeName(uint64_t code, std::string &nameOut);" << std::endl;
    outputHeader << "}" << std::endl;
    outputHeader << "}" << std::endl;

    std::ofstream outputSource(argv[2]);
    if (!outputSource.is_open()) {
        std::cerr << "Failed to open source output file " << argv[2] << std::endl;
        return 1;
    }

    outputSource << "// AUTO GENERATED FILE, DO NOT EDIT!" << std::endl;
    outputSource << "// This file has been generated by nstatus-gen from the systeminfo library" << std::endl;
    outputSource << std::endl;
    outputSource << "#include <unordered_map>" << std::endl;
    outputSource << "#include <cstdint>" << std::endl;
    outputSource << "#include <string>" << std::endl;
    outputSource << std::endl;
    outputSource << "namespace Sys {" << std::endl;
    outputSource << "namespace Win32 {" << std::endl;
    outputSource << "static std::unordered_map<uint64_t, std::string> NTSTATUS_CODES = {" << std::endl;

    bool first = true;
    for (const auto &status: codes) {
        if (first) {
            first = false;
        } else {
            outputSource << "," << std::endl;
        }

        outputSource << "    {0x" << std::hex << status.value << std::dec << ", \"" << status.name << "\"}";
    }

    outputSource << std::endl;
    outputSource << "};" << std::endl;
    outputSource << "bool lookupNtStatusCodeName(uint64_t code, std::string &nameOut) {" << std::endl;
    outputSource << "    auto it = NTSTATUS_CODES.find(code);" << std::endl;
    outputSource << "    if(it != NTSTATUS_CODES.end()) {" << std::endl;
    outputSource << "        nameOut = it->second;" << std::endl;
    outputSource << "        return true;" << std::endl;
    outputSource << "    }" << std::endl;
    outputSource << "    return false;" << std::endl;
    outputSource << "}" << std::endl;
    outputSource << "}" << std::endl;
    outputSource << "}" << std::endl;

    return 0;
}