#include <unistd.h>
#include <cstdlib>
#include <iostream>
#include <getopt.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#ifdef WITH_GUI
# include <wx/wx.h>
#endif

#include "m8/run.h"
#ifdef WITH_TUI
# include "tui/Tui.hpp"
#endif
#include "utils.h"
#include "globals.h"
#include "paleinfo.h"
#ifdef _WIN32
# include "usb/driver.h"
#endif
#include "usb/pongo_helper.h"

void print_credits() {
    printf(
        "::\n"
        ":: Palera1n beta " PALERAIN_VERSION "\n"
        "::\n"
        ":: (c) 2026\n"
        "::\n"
        ":: ========  Made by  =======>\n"
        ":: Made by: asdfugil, kok3shidoll, claration, mineek\n"
        ":: nekohaxx, plooshi, staturnz\n"
        ":: ======== Thanks to =======>\n"
        ":: Thanks to: itsnebulalol, llsc12, lrdsnow, dedbeddedbed\n"
        ":: kirb, ehilwyma, opa334, 0x7ff, alfiecg25, sneko, sbingner\n"
        ":: nikias, tihmstar\n"
        ":: Checkra1n (Siguza, axi0mx, littlelailo et al.)\n"
        ":: Procursus (Hayden Seay, Cameron Katri, Keto et al.)\n"
        ":: ==========================>\n\n"
    );
}

void print_usage(char* argv) {
    printf(
        "Usage: %s [OPTIONS]\n\n"
        "  -h, --help       Print usage information\n"
        "  -v, --version    Print version\n\n"
        "Options:\n\n"
        "  --cli                                 [f] Use command line interface\n"
        #ifdef WITH_GUI
        "  --gui                                 [f] Use graphical user interface\n"
        #endif
        #ifdef WITH_TUI
        "  --tui                                 [f] Use terminal user interface\n"
        #endif
        "  --dark-blockchain                     [f] Enable dark blockchain\n"
        "  --force-revert                        [f] Force environment reversion\n"
        "  -l, --rootless                        [f] Enable rootless mode (standard)\n"
        "  -f, --rootful                         [f] Boots fakefs\n"
        "  -c, --setup-fakefs                    [f] Setup fake filesystem\n"
        "  -B, --setup-partial-fakefs            [f] Setup partial fake filesystem\n"
        "  -s, --safe-mode                       [f] Enable safe mode\n"
        "  -T, --telnetd                         [f] Enable TELNET daemon on port 46 (insecure)\n"
        "  -V, --verbose-boot                    [f] Enable verbose booting\n"
        "  -p, --early-exit                      [f] Exit after uploading Pongo\n"
        "  -e, --extra-bootargs <BOOTARGS>       Set extra bootargs\n"
        "  -k, --override-pongo <FILE PATH>      Override Pongo image\n"
        "  -K, --override-kpf <FILE PATH>        Override kernel patchfinder\n"
        "  -o, --override-overlay <FILE PATH>    Override overlay\n"
        "  -r, --override-ramdisk <FILE PATH>    Override ramdisk\n"
        "  -n, --no-colors                       [f] Disable colors on the command line\n"
        "  -q, --quick                           [f] Enable Quick Mode\n"
        #ifdef _WIN32
        "  --install-drivers <PID>               Install libusbK drivers for PID\n"
        "  --remove-drivers <PID>                Remove libusbK drivers for PID\n"
        #endif
        , argv
    );

    printf(
        "\n"
        "Many things included in this program are made and developed by the Checkra1n team.\n"
        "Please do NOT contact the Checkra1n team for Palera1n, they are not responsible for this program.\n"
        "The original Checkra1n payloads are from: https://checkra.in/1337\n"
    );
}

void parse_arguments(int argc, char* argv[]) {
    int options;
    int option_index = 0;

    static struct option long_options[] = {
        {"help", no_argument, NULL, 'h'},
        {"version", no_argument, NULL, 'v'},
        {"cli", no_argument, NULL, 1},
        #ifdef WITH_GUI
        {"gui", no_argument, NULL, 2},
        #endif
        #ifdef WITH_GUI
        {"tui", no_argument, NULL, 3},
        #endif
        {"dark-blockchain", no_argument, NULL, 4},
        {"force-revert", no_argument, NULL, 5},
        {"rootless", no_argument, NULL, 'l'},
        {"rootful", no_argument, NULL, 'f'},
        {"setup-fakefs", no_argument, NULL, 'c'},
        {"setup-partial-fakefs", no_argument, NULL, 'B'},
        {"safe-mode", no_argument, NULL, 's'},
        {"telnetd", no_argument, NULL, 'T'},
        {"verbose-boot", no_argument, NULL, 'V'},
        {"early-exit", no_argument, NULL, 'p'},
        {"extra-bootargs", required_argument, NULL, 'e'},
        {"override-pongo", required_argument, NULL, 'k'},
        {"override-kpf", required_argument, NULL, 'K'},
        {"override-overlay", required_argument, NULL, 'o'},
        {"override-ramdisk", required_argument, NULL, 'r'},
        {"no-colors", no_argument, NULL, 'n'},
        {"quick", no_argument, NULL, 'q'},
        #ifdef _WIN32
        {"install-drivers", required_argument, NULL, 6},
        {"remove-drivers", required_argument, NULL, 7},
        #endif
        {NULL, 0, NULL, 0}
    };

    #ifdef WITH_GUI
    palerain_flags |= palerain_option_gui;
    # ifndef _WIN32
    if (isatty(STDIN_FILENO) && getenv("LLVM_PROFILE_FILE") == nullptr) {
        #ifdef WITH_TUI
        palerain_flags &= ~palerain_option_gui;
        palerain_flags |= palerain_option_tui;
        #else
        palerain_flags &= ~palerain_option_gui;
        palerain_flags |= palerain_option_cli;
        #endif
    } else {
        palerain_flags &= ~palerain_option_cli;
        palerain_flags |= palerain_option_gui;
        // xcode doesn't know how to display colors
        palerain_flags |= palerain_option_no_colors;
    }
    # endif
    #else
    palerain_flags |= palerain_option_cli;
    #endif

    while ((options = getopt_long(argc, argv, "hvlfcB:sTVpe:k:K:o:r:nq", long_options, &option_index)) != -1) {
        switch (options) {
            case 'h': // --help
                print_usage(argv[0]);
                exit(1);
            case 'v': // --version
                printf("Palera1n beta " PALERAIN_VERSION " [USB: %s (openra1n)]\n",
                    #ifdef __APPLE__
                    "IOKit"
                    #else
                    "libusb"
                    #endif
                );
                printf("%s (%s)\n", GIT_HASH, GIT_BRANCH);
                printf("Git: %s\n", GIT_URL);
                printf("Compiled on %s at %s\n", __DATE__, __TIME__);
                exit(0);
            case 1: // --cli
                palerain_flags &= ~palerain_option_tui;
                palerain_flags &= ~palerain_option_gui;
                palerain_flags |= palerain_option_cli;
                break;
            #ifdef WITH_GUI
            case 2: // --gui
                palerain_flags &= ~palerain_option_cli;
                palerain_flags &= ~palerain_option_tui;
                palerain_flags |= palerain_option_gui;
                break;
            #endif
            #ifdef WITH_TUI
            case 3: // --tui
                palerain_flags &= ~palerain_option_cli;
                palerain_flags &= ~palerain_option_gui;
                palerain_flags |= palerain_option_tui;
                break;
            #endif
            case 4: // --dark-blockchain
                palerain_flags |= palerain_option_flower_chain;
                break;
            case 5: // --force-revert
                palerain_flags |= palerain_option_force_revert;
                break;
            case 'l': // --rootless
                palerain_flags &= ~palerain_option_rootful;
                palerain_flags |= palerain_option_rootless;
                break;
            case 'f': // --rootful
                palerain_flags &= ~palerain_option_rootless;
                palerain_flags |= palerain_option_rootful;
                break;
            case 'c': // --setup-fakefs
                palerain_flags |= palerain_option_setup_rootful;
                break;
            case 'B': // --setup-partial-fakefs
                palerain_flags |= palerain_option_setup_partial_root;
                break;
            case 's': // --safe-mode
                palerain_flags |= palerain_option_safemode;
                break;
            case 'T': // --telnetd
                palerain_flags |= palerain_option_telnetd;
                break;
            case 'V': // --verbose-boot
                palerain_flags |= palerain_option_verbose_boot;
                break;
            case 'p': // --early-exit
                palerain_flags |= palerain_option_pongo_exit;
                break;
            case 'e': // --extra-bootargs
                if (strlen(optarg) > (sizeof(boot_args) - 0x20)) {
                    LOG_ERROR("Boot arguments too long");
                    exit(1);
                }
                snprintf(boot_args, sizeof(boot_args), "%s", optarg);
                break;
            case 'k': // --override-pongo
                if (!override_payload_from_file(optarg, &g_payload_pongo)) {
                    LOG_ERROR("Failed to load pongo payload\n");
                    exit(1);
                }
                if (g_payload_pongo.data_len > PONGO_MAX_SZ) {
                    LOG_ERROR("Pongo payload too large: %zu bytes (max %zu)", g_payload_pongo.data_len, PONGO_MAX_SZ);
                    exit(1);
                }
                LOG("Overriding pongo payload with %s", optarg);
                break;
            case 'K': // --override-kpf
                if (!override_payload_from_file(optarg, &g_payload_kpf)) {
                    LOG_ERROR("Failed to load kpf payload\n");
                    exit(1);
                }
                if (g_payload_kpf.data_len < 4
                    || (memcmp(g_payload_kpf.data, MACHO_MAGIC_64, 4) != 0
                    && memcmp(g_payload_kpf.data, MACHO_MAGIC_32, 4) != 0))
                {
                    LOG_ERROR("Invalid kpf payload, is it macho?");
                    exit(1);
                }
                LOG("Overriding kpf payload with %s", optarg);
                break;
            case 'o': // --override-overlay
                if (!override_payload_from_file(optarg, &g_payload_overlay)) {
                    LOG_ERROR("Failed to load overlay payload\n");
                    exit(1);
                }
                LOG("Overriding overlay payload with %s", optarg);
                break;
            case 'r': // --override-ramdisk
                if (!override_payload_from_file(optarg, &g_payload_ramdisk)) {
                    LOG_ERROR("Failed to load ramdisk payload\n");
                    exit(1);
                }
                LOG("Overriding ramdisk payload with %s", optarg);
                break;
            case 'n': // --no-colors
                palerain_flags |= palerain_option_no_colors;
                break;
            case 'q': // --quick
                palerain_flags |= palerain_option_quick;
                break;
            #ifdef _WIN32
            case 6: { // --install-drivers <PID>
                unsigned short parsed_pid = (unsigned short)strtol(optarg, NULL, 16);
                if (parsed_pid == 0) {
                    LOG_ERROR("Invalid PID value provided\n");
                    exit(1);
                }
                install_libusbk_target(0x05AC, parsed_pid);
                exit(0);
            }
            case 7: { // --remove-drivers <PID>
                unsigned short parsed_pid = (unsigned short)strtol(optarg, NULL, 16);
                if (parsed_pid == 0) {
                    LOG_ERROR("Invalid PID value provided\n");
                    exit(1);
                }
                uninstall_libusbk_target(0x05AC, parsed_pid);
                exit(0);
            }
            #endif
            case '?':
                LOG_ERROR("Unknown option\n");
                print_usage(argv[0]);
                exit(1);
            default: break;
        }
    }

    if (!(palerain_flags & palerain_option_gui) &&
        !(palerain_flags & palerain_option_tui))
    {

        if (!(palerain_flags & palerain_option_rootful) &&
            !(palerain_flags & palerain_option_rootless))
        {

            LOG_ERROR("You must specify either -l, --rootless or -f, --rootful.\n");
            print_usage(argv[0]);
            exit(1);
        }
    }
}

int main(int argc, char* argv[], char* envp[]) {
    print_credits();
    parse_arguments(argc, argv);
    LOG("palera1n_flags: %llu", palerain_flags);

    // communicating with libusb on linux needs root
    #ifdef __linux__
    if (geteuid() != 0) {
        LOG_WARN("You are not running as root, this may cause issues when exploiting");
    }
    #endif

    #ifdef WITH_GUI
    if (palerain_flags & palerain_option_gui) {
        wxEntryStart(argc, argv);
        wxTheApp->CallOnInit();
        wxTheApp->OnRun();
        wxTheApp->OnExit();
        wxEntryCleanup();
        return 0;
    } else if (palerain_flags & palerain_option_tui) {
        ui_run();
        return 0;
    } else {
        shared_t state{};
        return exploit(&state);
    }
    #else
    shared_t state{};
    return exploit(&state);
    #endif
}