#include <iostream>
#include <string>
#include <vector>
#include <unistd.h>
#include <cstring>
#include <sstream>

#include "version.h"

// Custom base64 encoding/decoding tables
static const std::string base64_chars =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/";

std::string base64_encode(const std::string &in) {
    std::string out;
    int val = 0, valb = -6;
    for (unsigned char c : in) {
        val = (val << 8) + c;
        valb += 8;
        while (valb >= 0) {
            out.push_back(base64_chars[(val >> valb) & 0x3F]);
            valb -= 6;
        }
    }
    if (valb > -6) out.push_back(base64_chars[((val << 8) >> (valb + 8)) & 0x3F]);
    while (out.size() % 4) out.push_back('=');
    return out;
}

std::string base64_decode(const std::string &in) {
    std::vector<int> T(256, -1);
    for (int i = 0; i < 64; i++) T[base64_chars[i]] = i;
    std::string out;
    int val = 0, valb = -8;
    for (unsigned char c : in) {
        if (T[c] == -1) break;
        val = (val << 6) + T[c];
        valb += 6;
        if (valb >= 0) {
            out.push_back(char((val >> valb) & 0xFF));
            valb -= 8;
        }
    }
    return out;
}

// Recursively decode and print if nested bb64 command is found
void recursive_print(const std::string &decoded) {
    std::cout << decoded << std::endl;
    size_t pos = decoded.find("bb64 ");
    if (pos != std::string::npos) {
        std::istringstream iss(decoded.substr(pos));
        std::string cmd, arg;
        iss >> cmd >> arg;
        if (cmd == "bb64" && !arg.empty()) {
            std::string nested = base64_decode(arg);
            recursive_print(nested);
        }
    }
}

int main(int argc, char *argv[]) {
    if (argc < 2) {
        std::cerr << "bb64 version " << VERSION << ", by J842." << std::endl;
        // heredoc for instructions
        std::cerr << R"(    

Usage:
    bb64 BASE64COMMAND          Decodes and runs the command
    bb64 -i BASE64COMMAND       Displays the decoded command
    bb64 -e COMMAND             Encodes the command and prints the result

)" << std::endl;
        return -1;
    }
    std::string mode;
    int argi = 1;
    if (std::string(argv[1]) == "-i" || std::string(argv[1]) == "-e") {
        mode = argv[1];
        argi = 2;
    }
    if (mode == "-e") {
        // Encode the rest of the arguments as a single command string
        std::ostringstream oss;
        for (int i = argi; i < argc; ++i) {
            if (i > argi) oss << " ";
            oss << argv[i];
        }
        std::string encoded = base64_encode(oss.str());
        std::cout << encoded << std::endl;
        return 0;
    } else if (mode == "-i") {
        if (argi >= argc) {
            std::cerr << "No BASE64COMMAND provided for -i" << std::endl;
            return -1;
        }
        std::string decoded = base64_decode(argv[argi]);
        recursive_print(decoded);
        return 0;
    } else {
        // Default: decode and run
        std::string decoded = base64_decode(argv[argi]);
        if (decoded.empty()) {
            std::cerr << "Failed to decode base64 command." << std::endl;
            return -1;
        }
        // Replace current process with bash -c "decoded"
        execlp("bash", "bash", "-c", decoded.c_str(), (char *)nullptr);
        // If execlp returns, there was an error
        std::cerr << "Failed to execute command." << std::endl;
        return -1;
    }
}