#include "cli.hpp"

#include <iostream>
#include <cstring>
#include <vector>

namespace Cli {

bool IsVerboseFlag(const char* arg) {
    return arg && std::strcmp(arg, "--verbose") == 0;
}

void PrintBanner(){
    std::cout << R"raw(
██████╗ ██╗   ██╗███╗   ██╗██╗     ██████╗  █████╗ ██████╗ ███████╗██████╗ 
██╔══██╗╚██╗ ██╔╝████╗  ██║██║     ██╔══██╗██╔══██╗██╔══██╗██╔════╝██╔══██╗
██║  ██║ ╚████╔╝ ██╔██╗ ██║██║     ██║  ██║███████║██║  ██║█████╗  ██████╔╝
██║  ██║  ╚██╔╝  ██║╚██╗██║██║     ██║  ██║██╔══██║██║  ██║██╔══╝  ██╔══██╗
██████╔╝   ██║   ██║ ╚████║███████╗██████╔╝██║  ██║██████╔╝███████╗██║  ██║
╚═════╝    ╚═╝   ╚═╝  ╚═══╝╚══════╝╚═════╝ ╚═╝  ╚═╝╚═════╝ ╚══════╝╚═╝  ╚═╝
)raw";
}

void PrintUsage() {
        std::cout << 
        "Dynloader loader\n\n"
        "LOCAL:\n"
        "  loader.exe <file.bin> [--verbose]\n"
        "  loader.exe --process <name|PID> <file.bin> [--verbose]\n"
        "  loader.exe --enc AES <file.bin> [--verbose]\n\n"
        "FILELESS:\n"
        "  loader.exe --fileless --source <host[:port]> [--verbose]\n"
        "  loader.exe --fileless --source <host> --enc AES [--verbose]\n"
        "  loader.exe --fileless --source <host> --process <name|PID> [--enc AES] [--verbose]\n\n"
        "SERVER:\n"
        "  loader.exe --server <file.bin> [--port 8080] [--verbose]          plain\n"
        "  loader.exe --server --enc AES <file.bin> [--port 8080] [--verbose]\n\n"
        "REMOTE Usage:\n"
        "   Usage --process <name | PID>\n\n"
        "DEBUGGING MODE:\n"
        "   Usage: --verbosE\n";
        "FUN Part:\n"
        "   Use --banner to see suprise\n"
        "   Usage - ./loader.exe --banner\n\n";
}

static bool IsFlag(const char* a) {
    return a && a[0] == '-';
}

static bool IsModeFlag(const char* a) {
    return a && (
        std::strcmp(a, "--fileless") == 0 ||
        std::strcmp(a, "--server") == 0 ||
        std::strcmp(a, "--enc") == 0 ||
        std::strcmp(a, "--selftest") == 0);
}

Options Parse(int argc, char* argv[]) {
    Options opt{};

    if (argc < 2) {
        opt.mode = Mode::Help;
        return opt;
    }

    for (int i = 1; i < argc; ++i) {
        if (IsVerboseFlag(argv[i])) {
            opt.verbose = true;
        }
    }

    if (std::strcmp(argv[1], "--selftest") == 0) {
        opt.mode = Mode::Normal;
        opt.inputPath = "__selftest__";
        return opt;
    }

    if (std::strcmp(argv[1], "--banner") == 0){
        opt.mode = Mode::Banner;
        return opt;
    }

    if (std::strcmp(argv[1], "--server") == 0) {
        opt.mode = Mode::Server;
        // Server uses its own verbose channel (not loader syscall dumps).
        if (opt.verbose) {
            opt.server_verbose = true;
            opt.verbose = false;
        }
        for (int i = 2; i < argc; ++i) {
            if (IsVerboseFlag(argv[i])) {
                opt.server_verbose = true;
                opt.verbose = false;
                continue;
            }
            if (std::strcmp(argv[i], "--port") == 0 && i + 1 < argc) {
                opt.serverPort = static_cast<std::uint16_t>(std::atoi(argv[++i]));
            } else if (std::strcmp(argv[i], "--enc") == 0 && i + 2 < argc &&
                       std::strcmp(argv[i + 1], "AES") == 0) {
                opt.useAes = true;
                i += 2;
                if (i < argc && !IsFlag(argv[i])) {
                    opt.inputPath = argv[i];
                }
            } else if (!IsFlag(argv[i])) {
                opt.inputPath = argv[i];
            }
        }
        return opt;
    }

    if (std::strcmp(argv[1], "--enc") == 0 && argc >= 4 && std::strcmp(argv[2], "AES") == 0) {
        opt.mode = Mode::Encrypt;
        opt.useAes = true;
        opt.inputPath = argv[3];
        return opt;
    }

    if (std::strcmp(argv[1], "--fileless") == 0) {
        opt.mode = Mode::Fileless;
        for (int i = 2; i < argc; ++i) {
            if (IsVerboseFlag(argv[i])) {

                continue;
            }
            if (std::strcmp(argv[i], "--source") == 0 && i + 1 < argc) {
                opt.sourceHost = argv[++i];
            } else if (std::strcmp(argv[i], "--enc") == 0 && i + 1 < argc &&
                       std::strcmp(argv[i + 1], "AES") == 0) {
                opt.useAes = true;
                ++i;
            } else if (std::strcmp(argv[i], "--process") == 0 && i + 1 < argc) {
                opt.processTarget = argv[++i];
                opt.remoteIngestion = true;
            }
        }
        if (opt.sourceHost.empty()) {
            opt.mode = Mode::Help;
        }
        return opt;
    }

    // Scan any order: --process <target> <file.bin> [--verbose]
    std::string binPath;
    for (int i = 1; i < argc; ++i) {
        if (IsVerboseFlag(argv[i])) {
            continue;
        }
        if (std::strcmp(argv[i], "--process") == 0 && i + 1 < argc) {
            opt.processTarget = argv[++i];
            opt.remoteIngestion = true;
            continue;
        }
        if (!IsFlag(argv[i]) && !IsModeFlag(argv[i])) {
            if (binPath.empty()) {
                binPath = argv[i];
            }
        }
    }

    if (!binPath.empty()) {
        opt.mode = Mode::Normal;
        opt.inputPath = binPath;
        return opt;
    }

    if (!IsFlag(argv[1])) {
        opt.mode = Mode::Normal;
        opt.inputPath = argv[1];
        return opt;
    }

    opt.mode = Mode::Help;
    return opt;
}

} // namespace Cli